|
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
|
import torch |
|
|
|
|
|
def init(): |
|
global model, tokenizer, generator |
|
model_name = "niruemon/llm-swp" |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16) |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
|
|
|
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto") |
|
|
|
|
|
def inference(inputs): |
|
global generator |
|
|
|
|
|
prompt = inputs.get("text", "") |
|
if not prompt: |
|
return {"error": "No input text provided."} |
|
|
|
|
|
try: |
|
result = generator(prompt, max_length=150, num_return_sequences=1) |
|
generated_text = result[0]["generated_text"] |
|
return {"generated_text": generated_text} |
|
except Exception as e: |
|
return {"error": str(e)} |
|
|