from typing import  Dict, List, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace, HuggingFaceEndpoint
import torch

class EndpointHandler:
    def __init__(self, path1="Cognute02/llama_3_1_8B_4bit"):
        # load model and processor from path
        quantization_config = BitsAndBytesConfig(
                                    load_in_4bit=True,
                                    bnb_4bit_quant_type="nf4",
                                    bnb_4bit_compute_dtype="float16",
                                    bnb_4bit_use_double_quant=True,
                                )
        
        self.model = AutoModelForCausalLM.from_pretrained(path1, quantization_config=quantization_config)
        self.tokenizer = AutoTokenizer.from_pretrained(path1)
        self.pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Args:
            data (:obj:):
                includes the deserialized image file as PIL.Image
        """
        inputs = data['inputs']
        outputs = self.pipe(inputs, max_new_tokens=256, do_sample=False, return_full_text=False,
                          temperature=0.25, top_p=0.25, top_k=10)
        
        return outputs