Update app.py

app.py

@@ -127,19 +127,19 @@ def load_and_prepare_local_dataset(file_path, text_column, label_column, test_si
     except Exception as e:
         raise Exception(f"Error loading dataset: {str(e)}")
 
-def preview_dataset(file_path, text_column, label_column):
+def preview_dataset(uploaded_file, text_column, label_column):
     """Preview a dataset file"""
     try:
-        if not file_path:
-            return "Please select a dataset file first."
+        if uploaded_file is None:
+            return "Please upload a dataset file first."
         
-        if not os.path.exists(file_path):
-            return f"❌ File not found: {file_path}"
+        # Get the file path from the uploaded file
+        file_path = uploaded_file.name if hasattr(uploaded_file, 'name') else uploaded_file
         
         df = pd.read_csv(file_path)
         
         preview_info = []
-        preview_info.append(f"📊 **Dataset Preview: {file_path}**")
+        preview_info.append(f"📊 **Dataset Preview: {os.path.basename(file_path)}**")
         preview_info.append(f"- **Total rows:** {len(df)}")
         preview_info.append(f"- **Columns:** {list(df.columns)}")
         preview_info.append("")
@@ -165,6 +165,8 @@ def preview_dataset(file_path, text_column, label_column):
         
     except Exception as e:
         return f"❌ Error previewing dataset: {str(e)}"
+
+def login_to_hf(token):
     """Login to Hugging Face"""
     global TOKEN
     TOKEN = token
@@ -308,7 +310,7 @@ def predict_csv(csv_file, model_path):
     except Exception as e:
         return f"❌ CSV processing failed: {str(e)}"
 
-def train_model(dataset_file, text_column, label_column, num_epochs, batch_size, 
+def train_model(uploaded_file, text_column, label_column, num_epochs, batch_size, 
                 learning_rate, hf_token, push_to_hub, username, model_name):
     """Start the model training process with local data"""
     global TRAINING_LOGS, MODEL_PATH
@@ -330,15 +332,18 @@ def train_model(dataset_file, text_column, label_column, num_epochs, batch_size,
     else:
         hub_model_id = None
     
-    # Validate dataset file
-    if not dataset_file or not os.path.exists(dataset_file):
-        TRAINING_LOGS.append(f"❌ Dataset file not found: {dataset_file}")
+    # Validate uploaded file
+    if uploaded_file is None:
+        TRAINING_LOGS.append("❌ Please upload a dataset file")
         yield "\n".join(TRAINING_LOGS)
         return
     
+    # Get the file path from the uploaded file
+    dataset_file = uploaded_file.name if hasattr(uploaded_file, 'name') else uploaded_file
+    
     try:
         # Load and prepare the dataset
-        TRAINING_LOGS.append(f"📊 Loading dataset from {dataset_file}...")
+        TRAINING_LOGS.append(f"📊 Loading dataset from uploaded file...")
         yield "\n".join(TRAINING_LOGS)
         
         dataset_dict, final_text_col, final_label_col = load_and_prepare_local_dataset(
@@ -477,20 +482,15 @@ with gr.Blocks(title="BERT Complaint Classifier") as app:
         # Training Tab
         with gr.TabItem("Train Model"):
             gr.Markdown("### Train a New Model with Local Data")
-            gr.Markdown("Select your local CSV file and configure training parameters")
+            gr.Markdown("Upload your CSV file and configure training parameters")
             
-            # Dataset selection and preview
+            # Dataset upload
             with gr.Row():
-                with gr.Column(scale=2):
-                    dataset_file = gr.Dropdown(
-                        label="Select Dataset File",
-                        choices=[f for f in os.listdir(".") if f.endswith(".csv")],
-                        value=LOCAL_DATA_FILES[0] if LOCAL_DATA_FILES[0] in os.listdir(".") else None,
-                        allow_custom_value=True
-                    )
-                    
-                with gr.Column(scale=1):
-                    refresh_btn = gr.Button("🔄 Refresh Files", size="sm")
+                dataset_file = gr.File(
+                    label="Upload Dataset (CSV)",
+                    file_types=[".csv"],
+                    type="filepath"
+                )
                     
             # Column configuration
             with gr.Row():
@@ -507,7 +507,7 @@ with gr.Blocks(title="BERT Complaint Classifier") as app:
             
             # Dataset preview
             preview_btn = gr.Button("📊 Preview Dataset", variant="secondary")
-            dataset_preview = gr.Markdown("Select a dataset file and click 'Preview Dataset' to see its structure.")
+            dataset_preview = gr.Markdown("Upload a dataset file and click 'Preview Dataset' to see its structure.")
             
             # Training parameters
             with gr.Row():
@@ -587,22 +587,20 @@ with gr.Blocks(title="BERT Complaint Classifier") as app:
             train_btn = gr.Button("Start Training", variant="primary")
             training_output = gr.Textbox(label="Training Progress", lines=10)
             
-            # Connect the buttons
-            post_train_push_btn.click(
-                push_to_hub_after_training,
-                inputs=[
-                    gr.Textbox(value=MODEL_PATH, visible=False),
-                    post_train_username,
-                    post_train_model_name,
-                    post_train_token
-                ],
-                outputs=post_train_status
+            # Connect the preview button
+            preview_btn.click(
+                preview_dataset,
+                inputs=[dataset_file, text_column, label_column],
+                outputs=dataset_preview
             )
             
+            # Connect the training button
             train_btn.click(
                 train_model,
                 inputs=[
                     dataset_file,
+                    text_column,
+                    label_column,
                     num_epochs,
                     batch_size,
                     learning_rate,
@@ -614,6 +612,18 @@ with gr.Blocks(title="BERT Complaint Classifier") as app:
                 outputs=training_output,
                 show_progress="full"
             )
+            
+            # Connect the post-training push button
+            post_train_push_btn.click(
+                push_to_hub_after_training,
+                inputs=[
+                    gr.Textbox(value=MODEL_PATH, visible=False),
+                    post_train_username,
+                    post_train_model_name,
+                    post_train_token
+                ],
+                outputs=post_train_status
+            )
 
         # Classification Tab
         with gr.TabItem("Classify Complaints"):