File size: 1,001 Bytes
ed323cf | 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 | 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}]
|