File size: 1,953 Bytes
74c671e 5b5d064 ae4c9ee 860868b 9c2e5c3 74c671e e1f255a |
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 51 52 53 54 55 56 57 58 59 60 |
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch
class EndpointHandler:
def __init__(self, path=""):
# Load tokenizer and model
self.tokenizer = T5Tokenizer.from_pretrained(path)
self.model = T5ForConditionalGeneration.from_pretrained(path)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
self.model.eval()
def __call__(self, data):
"""
Expected input JSON format:
{
"inputs": "Instruction: Generate the correct Frappe query for the given question...",
"parameters": {
"max_new_tokens": 128,
"temperature": 0.3,
"do_sample": false
}
}
"""
# Extract text and optional parameters
inputs = data.get("inputs", "")
params = data.get("parameters", {})
# Tokenize the prompt
encoded_input = self.tokenizer(
inputs,
return_tensors="pt",
padding=True,
truncation=True,
max_length=512
).to(self.device)
# Generate with optional parameters
gen_kwargs = {
"max_new_tokens": params.get("max_new_tokens", 128),
"temperature": params.get("temperature", 0.3),
"do_sample": params.get("do_sample", False),
"num_beams": params.get("num_beams", 1),
}
# Generate output
with torch.no_grad():
generated_ids = self.model.generate(**encoded_input, **gen_kwargs)
# Decode output
output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=False)
output=output.replace("<pad>","")
output=output.replace("</s>","")
output_text = output.replace("[BT]","`")
# Return as Hugging Face API expects
return [{"generated_text": output_text}]
|