| from typing import Dict, Any |
| import torch |
| from peft import PeftModel |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| BASE_MODEL_ID = "unsloth/DeepSeek-R1-Distill-Llama-70B-unsloth-bnb-4bit" |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID) |
|
|
| base_model = AutoModelForCausalLM.from_pretrained( |
| BASE_MODEL_ID, |
| torch_dtype=torch.bfloat16, |
| device_map="auto", |
| ) |
|
|
| self.model = PeftModel.from_pretrained(base_model, path) |
| self.model.eval() |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
| inputs = data.get("inputs", "") |
| parameters = data.get("parameters", {}) |
|
|
| input_ids = self.tokenizer( |
| inputs, return_tensors="pt" |
| ).input_ids.to(self.model.device) |
|
|
| with torch.no_grad(): |
| output_ids = self.model.generate( |
| input_ids, |
| max_new_tokens=parameters.get("max_new_tokens", 256), |
| temperature=parameters.get("temperature", 0.7), |
| ) |
|
|
| generated_text = self.tokenizer.decode( |
| output_ids[0], skip_special_tokens=True |
| ) |
|
|
| return {"generated_text": generated_text} |
