from typing import Dict, Any
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE_MODEL_ID = "unsloth/DeepSeek-R1-Distill-Llama-70B-unsloth-bnb-4bit"


class EndpointHandler:
    def __init__(self, path=""):
        self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)

        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL_ID,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )

        self.model = PeftModel.from_pretrained(base_model, path)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.get("inputs", "")
        parameters = data.get("parameters", {})

        input_ids = self.tokenizer(
            inputs, return_tensors="pt"
        ).input_ids.to(self.model.device)

        with torch.no_grad():
            output_ids = self.model.generate(
                input_ids,
                max_new_tokens=parameters.get("max_new_tokens", 256),
                temperature=parameters.get("temperature", 0.7),
            )

        generated_text = self.tokenizer.decode(
            output_ids[0], skip_special_tokens=True
        )

        return {"generated_text": generated_text}