File size: 1,720 Bytes
7c429d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | # handler.py
from typing import Dict, List, Any
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
class EndpointHandler:
def __init__(self, path: str = ""):
"""
Load model and tokenizer at startup.
`path` is the model repo on the Hub or local directory.
"""
# Determine device
self.device = 0 if torch.cuda.is_available() else -1
# Load tokenizer and model
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForCausalLM.from_pretrained(path).to(self.device)
# Set up a text-generation pipeline
self.generator = pipeline(
task="text-generation",
model=self.model,
tokenizer=self.tokenizer,
device=self.device
)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Handle each inference request. Must return a non-empty list.
`data` will always contain at least the key "inputs".
"""
# Extract input messages or fall back to raw prompt
raw = data.get("inputs", data)
prompt = ""
# Support both string inputs and OpenAI-style messages
if isinstance(raw, list):
for msg in raw:
role = msg.get("role", "user")
content = msg.get("content", "")
prompt += f"{role}: {content}\n"
else:
prompt = str(raw)
# Run generation (returns a list of dicts)
outputs = self.generator(prompt, max_new_tokens=128)
# Ensure we output a list of {"generated_text": ...}
return [{"response": out["response"]} for out in outputs]
|