Update pages/21_NLP_Transformer.py

pages/21_NLP_Transformer.py

@@ -29,6 +29,7 @@ def load_and_preprocess_data():
         return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
     encoded_dataset = dataset.map(preprocess_function, batched=True)
     encoded_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+    encoded_dataset = encoded_dataset.rename_column("label", "labels")  # Rename the column to 'labels'
     return DataLoader(encoded_dataset, shuffle=True, batch_size=batch_size)
 
 train_dataloader = load_and_preprocess_data()
@@ -48,7 +49,7 @@ if st.sidebar.button("Train"):
     for epoch in range(num_epochs):
         for batch in train_dataloader:
             batch = {k: v.to(device) for k, v in batch.items()}
-            outputs = model(**batch, labels=batch["labels"])  # Fixed input key
+            outputs = model(**batch)  # No need to pass labels explicitly if they are in the batch dictionary
             loss = outputs.loss
             loss.backward()