Update app.py

app.py

@@ -4,7 +4,9 @@ import time
 import os
 import pandas as pd
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
+from sklearn.metrics import confusion_matrix
+from sklearn.model_selection import train_test_split
 from tqdm import tqdm  # For progress bar during training
 
 # Set up Streamlit page
@@ -20,6 +22,10 @@ hardware = st.sidebar.selectbox("Hardware", ["CPU", "Single GPU", "Multi-GPU", "
 model_choice = st.sidebar.selectbox("Choose Model", ["bert-base-uncased", "distilbert-base-uncased", "roberta-base"])
 dataset_source = st.sidebar.selectbox("Dataset Source", ["glue/sst2", "imdb", "ag_news", "Custom"])
 
+# Column Mapping for custom datasets
+text_column = st.sidebar.text_input("Text Column", "text")
+label_column = st.sidebar.text_input("Label Column", "label")
+
 # Training Parameters
 epochs = st.sidebar.slider("Number of Epochs", 1, 10, 3)
 batch_size = st.sidebar.selectbox("Batch Size", [8, 16, 32, 64], index=1)
@@ -62,20 +68,36 @@ def train_model():
     model = AutoModelForSequenceClassification.from_pretrained(model_choice, num_labels=2)  # Adjust num_labels as necessary
 
     # Load dataset
-    dataset = load_dataset(dataset_source)
+    if dataset_source.lower() != "custom":
+        dataset = load_dataset(dataset_source)
+    else:
+        # Handle Custom Dataset
+        uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
+        if uploaded_file is not None:
+            dataset_df = pd.read_csv(uploaded_file)
+            dataset = Dataset.from_pandas(dataset_df)
 
     # Tokenization function
     def tokenize_function(examples):
-        return tokenizer(examples["text"], truncation=True, padding="max_length")
+        return tokenizer(examples[text_column], truncation=True, padding="max_length")
 
     tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-    # Ensure that the dataset has the correct label column (adjust the label column name if necessary)
-    if "label" not in tokenized_datasets["train"].features:
-        raise ValueError("Dataset does not have a 'label' column for supervised training")
-
-    train_dataset = tokenized_datasets["train"]
-    eval_dataset = tokenized_datasets["validation"]
+    # Handle missing or non-standard splits
+    if "train" in tokenized_datasets:
+        train_dataset = tokenized_datasets["train"]
+    else:
+        # Create a custom split if no train split exists
+        train_dataset = tokenized_datasets
+        train_dataset, eval_dataset = train_test_split(train_dataset, test_size=0.1)
+
+    # Check for validation or test split
+    if "validation" in tokenized_datasets:
+        eval_dataset = tokenized_datasets["validation"]
+    elif "test" in tokenized_datasets:
+        eval_dataset = tokenized_datasets["test"]
+    else:
+        raise ValueError("Dataset does not have a 'validation' or 'test' split.")
 
     # Checkpoint Handling
     if resume_training and os.path.exists(checkpoint_path):
@@ -107,6 +129,11 @@ def train_model():
 
     # Training Loop with Progress Bar
     metrics = []
+    loss_values = []  # To store loss values for plotting
+    accuracy_values = []  # To store accuracy values for plotting
+    all_preds = []  # To store predictions for confusion matrix
+    all_labels = []  # To store true labels for confusion matrix
+
     with open(log_file, "w") as log_file_handle:
         log_file_handle.write("Starting training...\n")
         log_file_handle.flush()
@@ -135,12 +162,49 @@ def train_model():
             metrics.append({"epoch": epoch+1, "loss": results["eval_loss"], "accuracy": results.get("eval_accuracy", 0)})
             pd.DataFrame(metrics).to_csv(metrics_file, index=False)
 
+            loss_values.append(results["eval_loss"])
+            accuracy_values.append(results.get("eval_accuracy", 0))
+
+            # Collect predictions and labels for confusion matrix
+            all_preds.extend(results.get("logits", []))
+            all_labels.extend(eval_dataset["label"])
+
             # Update logs & metrics in UI
             log_area.text(log_text)
             st.line_chart(pd.DataFrame(metrics).set_index("epoch"))
 
             time.sleep(2)
 
+    # After training, plot charts for loss, accuracy, and confusion matrix
+    plot_metrics(loss_values, accuracy_values)
+    plot_confusion_matrix(all_labels, all_preds)
+
+def plot_metrics(loss_values, accuracy_values):
+    # Plot Loss Curve using Streamlit chart
+    metrics_df = pd.DataFrame({
+        "Epoch": range(1, len(loss_values) + 1),
+        "Loss": loss_values,
+        "Accuracy": accuracy_values
+    })
+
+    st.write("### Training Loss and Accuracy Curve")
+    st.line_chart(metrics_df.set_index("Epoch"))
+
+def plot_confusion_matrix(true_labels, preds):
+    # Convert logits to predicted class labels
+    pred_labels = torch.argmax(torch.tensor(preds), axis=1).numpy()
+
+    # Compute confusion matrix
+    cm = confusion_matrix(true_labels, pred_labels)
+    
+    # Plot confusion matrix using Streamlit chart
+    fig, ax = plt.subplots(figsize=(8, 6))
+    ax = sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=["Class 0", "Class 1"], yticklabels=["Class 0", "Class 1"])
+    ax.set_title("Confusion Matrix")
+    ax.set_xlabel("Predicted Label")
+    ax.set_ylabel("True Label")
+    st.pyplot(fig)
+
 # Start Training
 if start_train:
     train_model()