File size: 1,025 Bytes
5548911 4d6f0eb aaf1302 96e6439 aaf1302 96e6439 aaf1302 5548911 aaf1302 5548911 7289a72 aaf1302 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | from transformers import AutoModelForCausalLM,AutoTokenizer,BitsAndBytesConfig
import torch
import os
class EndpointHandler():
def __init__(self, model_id=""):
self.device = "cuda:0"
self.bnb_config = BitsAndBytesConfig(load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,)
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.model = AutoModelForCausalLM.from_pretrained(model_id,
device_map={"":0},
quantization_config=self.bnb_config,)
def __call__(self, input:str) -> str:
inputs = self.tokenizer(input, return_tensors="pt").to(self.device)
outputs = self.model.generate(**inputs, max_new_tokens=20)
result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
return result |