File size: 2,054 Bytes
cc16cfc dba7cc7 cc16cfc dba7cc7 cc16cfc dba7cc7 cc16cfc dba7cc7 cc16cfc | 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 Dict, Any, List
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
class EndpointHandler:
def __init__(self, path=""):
base_id = "microsoft/Phi-3-mini-4k-instruct"
adapter_id = "CarlosMM24/phi3-mini-med-adapter"
self.tokenizer = AutoTokenizer.from_pretrained(base_id)
base = AutoModelForCausalLM.from_pretrained(
base_id,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto" if torch.cuda.is_available() else None,
)
self.model = PeftModel.from_pretrained(base, adapter_id)
self.model.eval()
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
prompt = data.get("inputs", data)
# 1) Apply chat template so Phi-3 knows this is a user turn and it should respond
messages = [{"role": "user", "content": prompt}]
formatted = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True, # IMPORTANT: adds the assistant turn marker
)
# 2) Tokenize and move to model device
inputs = self.tokenizer(formatted, return_tensors="pt")
device = next(self.model.parameters()).device
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
out = self.model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"], # IMPORTANT
max_new_tokens=64,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
)
# 3) Decode ONLY newly generated tokens (completion)
prompt_len = inputs["input_ids"].shape[1]
completion_ids = out[0, prompt_len:]
generated = self.tokenizer.decode(completion_ids, skip_special_tokens=True).strip()
return [{"generated_text": generated}] |