| 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} |
