| 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 |