| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig |
| import torch |
|
|
| class EndpointHandler: |
| def __init__(self): |
| model_id = "klusai/tf2-12b" |
| quant = BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_use_double_quant=True, |
| ) |
| self.tokenizer = AutoTokenizer.from_pretrained(model_id) |
| self.model = AutoModelForCausalLM.from_pretrained( |
| model_id, |
| device_map="auto", |
| quantization_config=quant, |
| torch_dtype=torch.bfloat16, |
| ) |
|
|
| def __call__(self, payload: dict): |
| prompt = payload.get("inputs") or payload.get("prompt", "") |
| max_new = payload.get("max_new_tokens", 1024) |
| inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device) |
| outputs = self.model.generate(**inputs, max_new_tokens=max_new) |
| return {"generated_text": self.tokenizer.decode(outputs[0], skip_special_tokens=True)} |
|
|
