from typing import Dict, Any, List
import random
from unsloth import FastLanguageModel
import torch

max_seq_length = 2048
dtype = None
load_in_4bit = True

class EndpointHandler:
    def __init__(self, path=""):
        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
            model_name = path, # YOUR MODEL YOU USED FOR TRAINING
            max_seq_length = max_seq_length,
            dtype = dtype,
            load_in_4bit = load_in_4bit
        )
        FastLanguageModel.for_inference(self.model) # Enable native 2x faster inference
        # pass

    def __call__(self, data: Dict[str, Any]) -> Any:
        inputs = self.tokenizer(data["inputs"], return_tensors = "pt").to("cuda")
        outputs = self.model.generate(**inputs, max_new_tokens = 64)

        prompt_length = inputs["input_ids"].shape[1]
        # Decode the generated output, skipping the tokens corresponding to the prompt
        text = self.tokenizer.decode(outputs[0][prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=True)
        return text