from typing import Dict, List, Any
from transformers import AutoModel
import numpy as np

class EndpointHandler():
    def __init__(self, path=""):
        self.model = AutoModel.from_pretrained("pabloruizponce/in2IN", trust_remote_code=True)

    def __call__(self, data: Dict[str, str]) -> Dict[str, np.ndarray]:

        interaction_text = data['interaction_text']
        individual1_text = data['individual1_text']
        individual2_text = data['individual2_text']
        
        prediction = self.model(interaction_text, individual1_text, individual2_text)
        prediction = {'individual1': prediction[0], 'individual2': prediction[1]}
        return prediction