File size: 1,196 Bytes
36284ba cd6f314 36284ba cd6f314 36284ba cd6f314 |
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 |
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
class EndpointHandler:
def __init__(self, model_dir):
# Load tokenizer and model during initialization
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForCausalLM.from_pretrained(model_dir)
def __call__(self, data):
"""
This method processes input data and generates output.
:param data: Input data, usually a dictionary with 'inputs' key.
"""
# Extract input prompt
inputs = data.get("inputs", "")
if not inputs:
return {"error": "No input provided"}
# Preprocess input
encoded_inputs = self.tokenizer(inputs, return_tensors="pt", padding=True)
# Generate output
with torch.no_grad():
outputs = self.model.generate(
**encoded_inputs,
max_length=200,
temperature=0.7,
do_sample=True
)
# Decode and return response
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
return {"generated_text": response}
|