from transformers import AutoModelForCausalLM, AutoTokenizer

class EndpointHandler:
    def __init__(self, path="Chituyi7/EBO-AlpacaLlama3-8B-InstructionTuned"):
        # Initialize your model here
        self.model = AutoModelForCausalLM.from_pretrained(path)
        self.tokenizer = AutoTokenizer.from_pretrained(path)

    def __call__(self, data):
        # Use your model to make predictions here
        inputs = self.tokenizer(data["inputs"], return_tensors="pt")
        outputs = self.model(**inputs)
        return outputs.logits.argmax(-1).item()