File size: 646 Bytes
38a9729 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
class EndpointHandler:
def __init__(self, model_path: str, task="text-generation"):
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
self.pipe = pipeline(task=task, model=self.model, tokenizer=self.tokenizer)
def __call__(self, inputs: dict) -> dict:
prompt = inputs.get("inputs", "")
params = inputs.get("parameters", {})
outputs = self.pipe(prompt, **params)
return outputs
|