Update pages/21_NLP_Transformer.py

pages/21_NLP_Transformer.py

@@ -128,36 +128,45 @@ scheduler = get_linear_schedule_with_warmup(
 loss_fn = torch.nn.CrossEntropyLoss().to(device)
 model = model.to(device)
 
-# Training loop
-for epoch in range(EPOCHS):
-    train_acc, train_loss = train_epoch(
-        model,
-        train_data_loader,
-        loss_fn,
-        optimizer,
-        device,
-        scheduler,
-        len(train_df)
-    )
+# Streamlit app
+st.title("Sentiment Analysis with BERT")
+st.write("""
+This application allows you to train a BERT model for sentiment analysis on the IMDb dataset. 
+You can input a movie review and the model will predict whether the sentiment is positive or negative.
+""")
+
+if st.button("Train Model"):
+    with st.spinner("Training the model..."):
+        # Training loop
+        for epoch in range(EPOCHS):
+            train_acc, train_loss = train_epoch(
+                model,
+                train_data_loader,
+                loss_fn,
+                optimizer,
+                device,
+                scheduler,
+                len(train_df)
+            )
 
-    print(f'Epoch {epoch + 1}/{EPOCHS}')
-    print(f'Train loss {train_loss} accuracy {train_acc}')
+            st.write(f'Epoch {epoch + 1}/{EPOCHS}')
+            st.write(f'Train loss {train_loss} accuracy {train_acc}')
 
-    val_acc, val_loss = eval_model(
-        model,
-        test_data_loader,
-        loss_fn,
-        device,
-        len(test_df)
-    )
+            val_acc, val_loss = eval_model(
+                model,
+                test_data_loader,
+                loss_fn,
+                device,
+                len(test_df)
+            )
 
-    print(f'Val loss {val_loss} accuracy {val_acc}')
+            st.write(f'Val loss {val_loss} accuracy {val_acc}')
 
-# Save the model
-model.save_pretrained('bert-sentiment-model')
-tokenizer.save_pretrained('bert-sentiment-model')
+        # Save the model
+        model.save_pretrained('bert-sentiment-model')
+        tokenizer.save_pretrained('bert-sentiment-model')
+    st.success("Model training complete!")
 
-# Streamlit app
 model = BertForSequenceClassification.from_pretrained('bert-sentiment-model')
 tokenizer = BertTokenizer.from_pretrained('bert-sentiment-model')
 model = model.eval()