Update pages/21_NLP_Transformer.py

pages/21_NLP_Transformer.py

@@ -34,6 +34,9 @@ def load_and_preprocess_data():
 
 train_dataloader = load_and_preprocess_data()
 
+# Initialize training status
+training_completed = st.sidebar.empty()
+
 # Training loop
 if st.sidebar.button("Train"):
     optimizer = AdamW(model.parameters(), lr=learning_rate)
@@ -49,7 +52,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)  # No need to pass labels explicitly if they are in the batch dictionary
+            outputs = model(**batch)
             loss = outputs.loss
             loss.backward()
 
@@ -59,7 +62,7 @@ if st.sidebar.button("Train"):
             progress_bar.update(1)
             loss_values.append(loss.item())
 
-    st.sidebar.success("Training completed")
+    training_completed.success("Training completed")
 
     # Plot loss values
     st.write("### Training Loss")