Adicionar camada densa no topo do modelo Huggingface BERT

Question

Quero adicionar uma camada densa no topo do transformador do Modelo BERT, gerando estados ocultos brutos e, em seguida, ajustar o modelo resultante. Especificamente, estou usando este modelo básico. Isso é o que o modelo deve fazer:

1. Codifique a frase (um vetor com 768 elementos para cada token da frase)
2. Mantenha apenas o primeiro vetor (relacionado ao primeiro token)
3. Adicione uma camada densa no topo deste vetor, para obter a transformação desejada

Até agora, codifiquei com sucesso as frases:

from sklearn.neural_network import MLPRegressor

import torch

from transformers import AutoModel, AutoTokenizer

# List of strings
sentences = [...]
# List of numbers
labels = [...]

tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-italian-xxl-cased")
model = AutoModel.from_pretrained("dbmdz/bert-base-italian-xxl-cased")

# 2D array, one line per sentence containing the embedding of the first token
encoded_sentences = torch.stack([model(**tokenizer(s, return_tensors='pt'))[0][0][0]
                                 for s in sentences]).detach().numpy()

regr = MLPRegressor()
regr.fit(encoded_sentences, labels)

Dessa forma, posso treinar uma rede neural alimentando-a com as frases codificadas. No entanto, essa abordagem claramente não ajusta o modelo base de BERT. Alguém pode me ajudar? Como posso construir um modelo (possivelmente em pytorch ou usando a biblioteca Huggingface) que pode ser totalmente ajustado?

Answer 1

Há duas maneiras de fazer isso: Como você está procurando ajustar o modelo para uma tarefa downstream semelhante à classificação, pode usar diretamente:

BertForSequenceClassificationaula. Executa o ajuste fino da camada de regressão logística na dimensão de saída de 768.

Como alternativa, você pode definir um módulo personalizado, que criou um modelo bert com base nos pesos pré-treinados e adiciona camadas sobre ele.

from transformers import BertModel
class CustomBERTModel(nn.Module):
    def __init__(self):
          super(CustomBERTModel, self).__init__()
          self.bert = BertModel.from_pretrained("dbmdz/bert-base-italian-xxl-cased")
          ### New layers:
          self.linear1 = nn.Linear(768, 256)
          self.linear2 = nn.Linear(256, 3) ## 3 is the number of classes in this example

    def forward(self, ids, mask):
          sequence_output, pooled_output = self.bert(
               ids, 
               attention_mask=mask)

          # sequence_output has the following shape: (batch_size, sequence_length, 768)
          linear1_output = self.linear1(sequence_output[:,0,:].view(-1,768)) ## extract the 1st token's embeddings

          linear2_output = self.linear2(linear2_output)

          return linear2_output

tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-italian-xxl-cased")
model = CustomBERTModel() # You can pass the parameters if required to have more flexible model
model.to(torch.device("cpu")) ## can be gpu
criterion = nn.CrossEntropyLoss() ## If required define your own criterion
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()))

for epoch in epochs:
    for batch in data_loader: ## If you have a DataLoader()  object to get the data.

        data = batch[0]
        targets = batch[1] ## assuming that data loader returns a tuple of data and its targets
        
        optimizer.zero_grad()   
        encoding = tokenizer.batch_encode_plus(data, return_tensors='pt', padding=True, truncation=True,max_length=50, add_special_tokens = True)
        outputs = model(input_ids, attention_mask=attention_mask)
        outputs = F.log_softmax(outputs, dim=1)
        input_ids = encoding['input_ids']
        attention_mask = encoding['attention_mask']
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        

Answer 2

Se você quiser ajustar o próprio modelo de BERT, você precisará modificar os parâmetros do modelo. Para fazer isso, você provavelmente desejará fazer seu trabalho com o PyTorch. Aqui está algum código psuedo aproximado para ilustrar:

from torch.optim import SGD

model = ... # whatever model you are using
parameters = model.parameters() # or some more specific set of parameters
optimizer = SGD(parameters,lr=.01) # or whatever optimizer you want
optimizer.zero_grad() # boiler-platy pytorch function

input = ... # whatever the appropriate input for your task is
label = ... # whatever the appropriate label for your task is
loss = model(**input, label) # usuall loss is the first item returned
loss.backward() # calculates gradient
optim.step() # runs optimization algorithm

Eu deixei de fora todos os detalhes relevantes porque eles são muito entediantes e específicos para qualquer tarefa específica. Huggingface tem um bom artigo explicando isso com mais detalhes aqui , e você definitivamente vai querer consultar alguma documentação do pytorch ao usar qualquer coisa do pytorch. Eu recomendo fortemente o pytorch blitz antes de tentar fazer algo sério com ele.