from typing import Dict, List, Any
import torch
from transformers import PreTrainedTokenizerFast, AutoModelForCausalLM, pipeline

class EndpointHandler():
    def __init__(self, path=""):
        # 1. Force the Fast Tokenizer to bypass the legacy TokenizersBackend bug
        # Setting padding_side="left" is standard for Causal LMs
        self.tokenizer = PreTrainedTokenizerFast.from_pretrained(
            path, 
            padding_side="left"
        )
        
        # Safely set pad token to eos_token to avoid warnings
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

        # 2. Load Model strictly for CPU
        # Using bfloat16 optimizes RAM usage significantly on Sapphire Rapids
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            torch_dtype=torch.bfloat16
        )
        
        # 3. Create Pipeline attached to self
        self.pipeline = pipeline(
            "text-generation", 
            model=self.model, 
            tokenizer=self.tokenizer,
            return_full_text=False # Prevent echoing the prompt
        )

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # Extract inputs and parameters safely
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", {})

        # Execute generation
        prediction = self.pipeline(inputs, **parameters)
        return prediction