Spaces:
Runtime error
Runtime error
File size: 4,093 Bytes
f2a699d 049e137 f2a699d 049e137 f2a699d 049e137 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
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 |