File size: 1,329 Bytes
012a858 520e9a8 d3cb92b 520e9a8 d3cb92b 520e9a8 012a858 96bf2b2 d3cb92b 520e9a8 d3cb92b 012a858 520e9a8 012a858 520e9a8 012a858 520e9a8 012a858 96bf2b2 | 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 | from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import os
class EndpointHandler:
def __init__(self, path=""):
# Explicitly prevent sentence-transformers auto-detection
os.environ["TRANSFORMERS_OFFLINE"] = "1"
print(f"Loading T5Gemma model from: {path}")
self.tokenizer = AutoTokenizer.from_pretrained(
path,
trust_remote_code=True
)
self.model = AutoModelForSeq2SeqLM.from_pretrained(
path,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
device_map="auto"
)
print("T5Gemma model loaded successfully")
def __call__(self, data):
inputs = data.pop("inputs", data)
messages = [{"role": "user", "content": inputs}]
input_ids = self.tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
return_tensors="pt"
)
outputs = self.model.generate(
input_ids,
max_new_tokens=1024,
temperature=0.1,
do_sample=True
)
return {
"generated_text": self.tokenizer.decode(outputs[0], skip_special_tokens=True)
}
|