Update app.py

app.py

@@ -4,8 +4,8 @@ import time
 import os
 import pandas as pd
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
-from datasets import load_dataset, Dataset, DatasetDict
-from sklearn.model_selection import train_test_split
+from datasets import load_dataset
+from tqdm import tqdm  # For progress bar during training
 
 # Set up Streamlit page
 st.set_page_config(page_title="AutoTrain AI", page_icon="🚀", layout="wide")
@@ -18,19 +18,7 @@ hf_user = st.sidebar.selectbox("Hugging Face User", ["hennings1984"])
 task = st.sidebar.selectbox("Select Task", ["Text Classification", "Sentiment Analysis"])
 hardware = st.sidebar.selectbox("Hardware", ["CPU", "Single GPU", "Multi-GPU", "TPU"])
 model_choice = st.sidebar.selectbox("Choose Model", ["bert-base-uncased", "distilbert-base-uncased", "roberta-base"])
-
-# Dataset Configuration
-dataset_source = st.sidebar.selectbox("Dataset Source", ["Hugging Face", "Upload Your Dataset"])
-
-if dataset_source == "Hugging Face":
-    # Choose Hugging Face dataset
-    dataset_name = st.sidebar.text_input("Enter Hugging Face Dataset Name", "imdb")
-else:
-    # Upload Custom Dataset
-    uploaded_file = st.sidebar.file_uploader("Upload Your Dataset (CSV/TSV/JSON)", type=["csv", "json", "tsv"])
-    # Allow the user to map columns
-    text_column = st.sidebar.text_input("Text Column Name", "text")
-    label_column = st.sidebar.text_input("Label Column Name", "label")
+dataset_source = st.sidebar.selectbox("Dataset Source", ["glue/sst2", "imdb", "ag_news", "Custom"])
 
 # Training Parameters
 epochs = st.sidebar.slider("Number of Epochs", 1, 10, 3)
@@ -63,44 +51,7 @@ log_area = st.empty()
 
 # Live Training Metrics
 st.write("### Training Metrics 📊")
-
-# Dataset Loading Logic
-def load_custom_dataset(uploaded_file, text_column, label_column):
-    if uploaded_file is not None:
-        file_type = uploaded_file.name.split('.')[-1]
-        if file_type == "csv":
-            df = pd.read_csv(uploaded_file)
-        elif file_type == "json":
-            df = pd.read_json(uploaded_file)
-        elif file_type == "tsv":
-            df = pd.read_csv(uploaded_file, sep="\t")
-        else:
-            st.error("Unsupported file type")
-            return None
-
-        # Map the columns based on user input
-        df = df[[text_column, label_column]]
-        return df
-    return None
-
-def load_huggingface_dataset(dataset_name):
-    return load_dataset(dataset_name)
-
-def split_dataset(df, text_column, label_column):
-    # Split the dataset into train, validation, and test (80/10/10 split)
-    train_df, temp_df = train_test_split(df, test_size=0.2, random_state=42)
-    val_df, test_df = train_test_split(temp_df, test_size=0.5, random_state=42)
-
-    # Convert to Hugging Face Dataset format
-    train_dataset = Dataset.from_pandas(train_df)
-    val_dataset = Dataset.from_pandas(val_df)
-    test_dataset = Dataset.from_pandas(test_df)
-
-    return DatasetDict({
-        "train": train_dataset,
-        "validation": val_dataset,
-        "test": test_dataset
-    })
+progress_bar = st.progress(0)  # Initialize progress bar
 
 # Training Function
 def train_model():
@@ -108,26 +59,21 @@ def train_model():
 
     # Load model & tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_choice)
-    model = AutoModelForSequenceClassification.from_pretrained(model_choice, num_labels=2)
+    model = AutoModelForSequenceClassification.from_pretrained(model_choice, num_labels=2)  # Adjust num_labels as necessary
 
     # Load dataset
-    if dataset_source == "Hugging Face":
-        dataset = load_huggingface_dataset(dataset_name)
-    else:
-        dataset_df = load_custom_dataset(uploaded_file, text_column, label_column)
-        if dataset_df is not None:
-            # Split dataset if it's not already split
-            dataset = split_dataset(dataset_df, text_column, label_column)
-
-    # Check the dataset structure to identify the correct column name
-    st.write(f"Dataset columns: {dataset['train'].column_names}")
-    
+    dataset = load_dataset(dataset_source)
+
     # Tokenization function
     def tokenize_function(examples):
-        # Adjust this based on the actual column name
-        return tokenizer(examples[text_column], truncation=True, padding="max_length")
+        return tokenizer(examples["text"], 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"]
 
@@ -159,14 +105,22 @@ def train_model():
         eval_dataset=eval_dataset,
     )
 
-    # Training Loop
+    # Training Loop with Progress Bar
     metrics = []
     with open(log_file, "w") as log_file_handle:
         log_file_handle.write("Starting training...\n")
         log_file_handle.flush()
 
         for epoch in range(epochs):
-            trainer.train()
+            # Initialize progress bar for this epoch
+            progress_bar.progress(0)  # Reset progress bar at the start of each epoch
+            
+            # Training with tqdm for real-time progress bar
+            for step, batch in enumerate(trainer.get_train_dataloader()):
+                trainer.training_step(model, batch)  # Perform a training step
+                progress_bar.progress((step + 1) / len(trainer.get_train_dataloader()))  # Update progress bar
+
+            # Evaluate the model at the end of each epoch
             results = trainer.evaluate()
 
             # Save Checkpoint