from typing import Dict, List, Any
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

class EndpointHandler:
    def __init__(self, path: str = ""):
        # Device to use for inference. "cuda" if CUDA is available, else "cpu"
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForMaskedLM.from_pretrained(path).to(self.device)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        data args:
            inputs (:obj: `str`)
        Return:
            A :obj:`list` | `dict`: will be serialized and returned
        """
        # Get inputs
        inputs = data.pop("inputs", data)
        if isinstance(inputs, str):
            inputs = [inputs]

        # Run the model to get logits
        inputs = self.tokenizer(
            inputs,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.model.config.max_position_embeddings,  # Max length according to the model configuration
        ).to(self.device)
        with torch.no_grad():
            logits = self.model(**inputs).logits

        # Apply transformations to get the sparse vectors
        inter = torch.log1p(torch.relu(logits))
        token_max = torch.max(inter, dim=1)

        nz_tokens_i, nz_tokens_j = torch.where(token_max.values > 0)

        output = []
        for i in range(token_max.values.shape[0]):
            nz_tokens = nz_tokens_j[nz_tokens_i == i]
            nz_weights = token_max.values[i, nz_tokens]
            output.append(
                {"indices": nz_tokens.tolist(), "values": nz_weights.tolist()}
            )

        return output[0] if isinstance(inputs, str) else output