vrutkovs/Lusterka-7B

1from typing import Any2 3import torch4from transformers import (5    AutoModelForCausalLM,6    AutoTokenizer,7    BitsAndBytesConfig,8    pipeline,9)10 11 12class EndpointHandler:13    def __init__(self, path: str = ""):14        self.tokenizer = AutoTokenizer.from_pretrained(path)15        self.model = AutoModelForCausalLM.from_pretrained(16            path,17            quantization_config=BitsAndBytesConfig(load_in_8bit=True),18            device_map="auto",19        )20        self.pipe = pipeline(21            "text-generation",22            model=self.model,23            tokenizer=self.tokenizer,24        )25 26    def __call__(self, data: dict[str, Any]) -> list[dict[str, Any]]:27        inputs = data.pop("inputs", data)28        parameters = data.pop("parameters", {})29 30        gen_kwargs = {31            "max_new_tokens": parameters.get("max_new_tokens", 256),32            "temperature": parameters.get("temperature", 0.8),33            "repetition_penalty": parameters.get("repetition_penalty", 1.3),34            "no_repeat_ngram_size": parameters.get("no_repeat_ngram_size", 4),35            "do_sample": parameters.get("do_sample", True),36        }37 38        outputs = self.pipe(inputs, **gen_kwargs)39        return outputs40