| from transformers import AutoTokenizer, AutoModelForCausalLM |
| from peft import PeftModel |
| import torch |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| base_model = "unsloth/mistral-7b-instruct-v0.2-bnb-4bit" |
| lora_path = path |
| self.tokenizer = AutoTokenizer.from_pretrained(base_model) |
| model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto", torch_dtype=torch.float16) |
| self.model = PeftModel.from_pretrained(model, lora_path) |
| self.model.eval() |
|
|
| def __call__(self, data): |
| inputs = data.get("inputs", data) |
| if isinstance(inputs, list): |
| prompt = inputs[0] |
| else: |
| prompt = inputs |
| inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device) |
| with torch.no_grad(): |
| outputs = self.model.generate(**inputs, max_new_tokens=300) |
| out = self.tokenizer.decode(outputs[0], skip_special_tokens=True) |
| return [{"generated_text": out}] |
|
|
