Update app.py

app.py

@@ -4,7 +4,8 @@ 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, DatasetDict
+from sklearn.model_selection import train_test_split
 
 # Set up Streamlit page
 st.set_page_config(page_title="AutoTrain AI", page_icon="🚀", layout="wide")
@@ -17,7 +18,19 @@ 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_source = st.sidebar.selectbox("Dataset Source", ["glue/sst2", "imdb", "ag_news", "Custom"])
+
+# 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")
 
 # Training Parameters
 epochs = st.sidebar.slider("Number of Epochs", 1, 10, 3)
@@ -51,6 +64,44 @@ 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
+    })
+
 # Training Function
 def train_model():
     st.success(f"Training started for {task} with {model_choice} on {device.upper()}")
@@ -60,11 +111,21 @@ def train_model():
     model = AutoModelForSequenceClassification.from_pretrained(model_choice, num_labels=2)
 
     # Load dataset
-    dataset = load_dataset(dataset_source)
-
+    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}")
+    
     # Tokenization function
     def tokenize_function(examples):
-        return tokenizer(examples["sentence"], truncation=True, padding="max_length")
+        # Adjust this based on the actual column name
+        return tokenizer(examples[text_column], truncation=True, padding="max_length")
 
     tokenized_datasets = dataset.map(tokenize_function, batched=True)
     train_dataset = tokenized_datasets["train"]