Spaces:
Sleeping
Sleeping
| 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 |