from torch import nn
from transformers import RobertaModel, RobertaConfig
import torch
from transformers import RobertaTokenizer

tokenizer = RobertaTokenizer.from_pretrained('roberta-base')

### Définition des tokens spéciaux ###
# indiquent respectivement le début de séquence, la séparation entre hypothèse et prémisse, et la fin de séquence #
bos_token = tokenizer.bos_token
sep_token = tokenizer.sep_token
eos_token = tokenizer.eos_token

class RobertaSNLI(nn.Module):
    def __init__(self):
        super(RobertaSNLI, self).__init__()
        config = RobertaConfig.from_pretrained('roberta-base')
        config.output_attentions = True  # activer sortie des poids d'attention
        config.max_position_embeddings = 130  # gérer la longueur des séquences
        config.hidden_size = 256    # taille des états cachés du modèle
        config.num_hidden_layers = 4    # nombre de couches cachées dans le transformateur
        config.intermediate_size = 512  # taille couche intermédiaire dans modèle de transformateur
        config.num_attention_heads = 4  # nombre de têtes d'attentions 

        self.roberta = RobertaModel(config)
        self.roberta.requires_grad = True
        self.output = nn.Linear(256, 3) # couche de sortie linéaire. Entrée la taille des états cachées et 3 sorties
        
    def forward(self, input_ids, attention_mask=None):
        outputs = self.roberta(input_ids, attention_mask=attention_mask)
        roberta_out = outputs[0]  # séquence des états cachés à la sortie de la dernière couche
        attentions = outputs.attentions   # poids d'attention du modèle RoBERTa
        
        return self.output(roberta_out[:, 0]), attentions
    

class LstmSNLI(nn.Module):
    def __init__(self):
        super(LstmSNLI, self).__init__()
        # Couche d'embedding : transforme les indices de mots en vecteurs denses
        self.embedding = nn.Embedding(
            num_embeddings=tokenizer.vocab_size,  # Taille du vocabulaire obtenu depuis le tokenizer
            embedding_dim=128,  # Dimension des vecteurs d'embedding
            padding_idx=tokenizer.pad_token_id,  # Index du token de padding utilisé pour égaliser les longueurs des séquences
            max_norm=1  # La norme maximale des vecteurs d'embedding; réduit la variance des embeddings
        )
        # Couche LSTM : réseau de neurones récurrent capable de capturer des dépendances à long terme
        self.lstm = nn.LSTM(
            num_layers=4,  # Utilisation de 4 couches LSTM empilées pour plus de profondeur
            input_size=128,  # Taille des entrées correspondant à la dimension des embeddings
            hidden_size=128,  # Taille des états cachés dans le LSTM
            batch_first=True,  # Indique que la première dimension des entrées représente la taille du batch
            dropout=0.1,  # Taux de dropout appliqué aux sorties de chaque couche LSTM, sauf la dernière
            bidirectional=True  # LSTM bidirectionnel pour capturer des informations contextuelles des deux directions
        )
        # Couche linéaire : effectue une transformation linéaire pour classifier les sorties du LSTM
        self.linear = nn.Linear(
            in_features=256,  # Taille d'entrée
            out_features=3  # Trois sorties pour les trois classes du SNLI
        )
        
    def forward(self, input_ids):
        # Transformation des indices de mots en embeddings
        embed = self.embedding(input_ids)
        # Passage des embeddings à travers le LSTM
        lstm_out = self.lstm(embed)  # lstm_out contient les sorties de toutes les étapes temporelles; les états cachés ne sont pas utilisés ici
        return self.linear(lstm_out[0][:,0])


def load_custom_model(model_path, model_type='lstm'):
    if model_type == 'lstm':
        model = LstmSNLI() 
    else:
        model = RobertaSNLI() 
        print("")
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))      # chargement des poids du modèle depuis le fichier
    model.eval()
    return model