File size: 1,823 Bytes
f0dcbdb | 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 | from typing import Any, Dict
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
class EndpointHandler:
def __init__(self, path=""):
base_model = AutoModelForCausalLM.from_pretrained(
"microsoft/Phi-3-mini-4k-instruct",
device_map="auto",
torch_dtype="auto",
)
self.model = PeftModel.from_pretrained(base_model, path)
self.tokenizer = AutoTokenizer.from_pretrained(
"microsoft/Phi-3-mini-4k-instruct"
)
self.model.eval()
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
question = data.get("inputs")
if not isinstance(question, str) or not question.strip():
raise ValueError("Envía una pregunta de texto en el campo 'inputs'.")
parameters = data.get("parameters") or {}
max_new_tokens = int(parameters.get("max_new_tokens", 256))
max_new_tokens = max(1, min(max_new_tokens, 512))
messages = [{"role": "user", "content": question}]
input_ids = self.tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt",
)
input_ids = input_ids.to(self.model.get_input_embeddings().weight.device)
with torch.inference_mode():
generated = self.model.generate(
input_ids=input_ids,
max_new_tokens=max_new_tokens,
do_sample=False,
eos_token_id=self.tokenizer.eos_token_id,
)
new_tokens = generated[0, input_ids.shape[-1]:]
answer = self.tokenizer.decode(
new_tokens,
skip_special_tokens=True,
).strip()
return {"generated_text": answer} |