18 lines
640 B
Python
18 lines
640 B
Python
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
def model_fn(model_dir):
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
|
model = AutoModelForCausalLM.from_pretrained(model_dir)
|
|
return {"model": model, "tokenizer": tokenizer}
|
|
|
|
def predict_fn(data, model_dict):
|
|
model = model_dict["model"]
|
|
tokenizer = model_dict["tokenizer"]
|
|
|
|
inputs = tokenizer(data["inputs"], return_tensors="pt").to(model.device)
|
|
outputs = model.generate(**inputs, max_length=128)
|
|
return tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
def output_fn(prediction, accept):
|
|
return {"generated_text": prediction}
|