from typing import Dict, List, Any
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline 

class EndpointHandler():
    def __init__(self, path=""):
      # Load model directly
        model = AutoModelForCausalLM.from_pretrained( 
            "jdgalvan/Phi-3-mini-128k-instruct",  
            device_map="cuda",  
            torch_dtype="auto",  
            trust_remote_code=True,  
        ) 
        
        tokenizer = AutoTokenizer.from_pretrained("jdgalvan/Phi-3-mini-128k-instruct") 

        self.pipe = pipeline( 
            "text-generation", 
            model=model, 
            tokenizer=tokenizer, 
        )

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None)

        # pass inputs with all kwargs in data
        if parameters is not None:
            prediction = self.pipe(inputs, **parameters)
        else:
            prediction = self.pipe(inputs)
            
        return prediction