Spaces:
Runtime error
Runtime error
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 |