from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline
import torch
import json

class EndpointHandler:
    def __init__(self, path=""):
        # Load tokenizer and model
        self.tokenizer = GPT2Tokenizer.from_pretrained(path)
        self.model = GPT2LMHeadModel.from_pretrained(
            path,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        
        # Set pad token if not set
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
    
    def __call__(self, data):
        """
        Args:
            data (dict): The payload sent to the endpoint
        Returns:
            dict: prediction response
        """
        inputs = data.pop("inputs", "")
        parameters = data.pop("parameters", {})
        
        # Set default parameters
        max_length = parameters.get("max_length", 100)
        temperature = parameters.get("temperature", 0.7)
        do_sample = parameters.get("do_sample", True)
        top_p = parameters.get("top_p", 0.9)
        
        # Tokenize input
        input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
        
        # Generate
        with torch.no_grad():
            outputs = self.model.generate(
                input_ids,
                max_length=max_length,
                temperature=temperature,
                do_sample=do_sample,
                top_p=top_p,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                attention_mask=torch.ones_like(input_ids)
            )
        
        # Decode output
        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return {"generated_text": generated_text}