Fine tunning t5: Too many values to unpack (expected 2)

🤗Transformers
1

I am trying to create a simple prof-of-concept of fine tunning t5 for generated_question:

import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
from torch.utils.data import DataLoader, Dataset

# Verifique se a GPU está disponível e use-a, se possível
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Carregue o modelo T5 pré-treinado e o tokenizador
model_name = 'unicamp-dl/ptt5-base-portuguese-vocab'
model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
tokenizer = T5Tokenizer.from_pretrained(model_name)

# Suponha que você tenha um conjunto de dados em que cada exemplo é um par de declaração e pergunta.
class QADataset(Dataset):
    def __init__(self, data, tokenizer, max_length=128):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        item = self.data[index]
        source_text = item['statement']
        target_text = item['question']

        source_tokens = self.tokenizer.encode(source_text, max_length=self.max_length, truncation=True, padding='max_length', return_tensors='pt')
        target_tokens = self.tokenizer.encode(target_text, max_length=self.max_length, truncation=True, padding='max_length', return_tensors='pt')

        return {
            'input_ids': source_tokens.to(device),
            'attention_mask': source_tokens.to(device) != self.tokenizer.pad_token_id,
            'labels': target_tokens.to(device),
        }

# Carregue seus dados
data = [
    {'statement': 'O Sol é uma estrela.', 'question': 'O que é o Sol?'},
    # Adicione mais exemplos aqui
]

# Crie um DataLoader com seu conjunto de dados
dataset = QADataset(data, tokenizer)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

# Defina os hiperparâmetros de treinamento
epochs = 3
learning_rate = 1e-4

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Treinamento
for epoch in range(epochs):
    model.train()
    # Dentro do loop de treinamento
    for batch in dataloader:
        input_ids = batch['input_ids']
        labels = batch['labels']
        # Certifique-se de que 'attention_mask' seja criado internamente pelo modelo
        optimizer.zero_grad()
        # Altere a chamada do modelo para ajustar 'input_ids' e 'labels'
        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()


# Após o treinamento, você pode usar o modelo para gerar perguntas
model.eval()
statement = "A Lua orbita a Terra."
input_ids = tokenizer.encode(statement, return_tensors='pt').to(device)
output = model.generate(input_ids, max_length=32, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95)
generated_question = tokenizer.decode(output[0], skip_special_tokens=True)

print("Declaração:", statement)
print("Pergunta gerada:", generated_question)

And got this error:

ValueError: too many values to unpack (expected 2)

from the line:

outputs = model(input_ids=input_ids, labels=labels)

Anyone knows how to fix?