Update app.py

app.py

@@ -231,7 +231,7 @@ def compute_metrics(eval_pred):
     }
 
 def train_model_inline(uploaded_file, text_column, label_column, num_epochs, batch_size, 
-                      learning_rate, hf_token, push_to_hub, username, model_name):
+                       learning_rate, hf_token, push_to_hub, username, model_name):
     """Train the model using inline training (no subprocess)"""
     global TRAINING_LOGS, MODEL_PATH, CURRENT_MODEL, CURRENT_TOKENIZER
     
@@ -642,22 +642,123 @@ def get_available_datasets():
         available_files = ["No CSV files found in current directory"]
     
     return available_files
+
 def display_available_datasets():
-                datasets = get_available_datasets()
-                if datasets:
-                    return "**Available CSV files:**\n\n" + "\n".join([f"- {file}" for file in datasets])
-                else:
-                    return "No CSV files found in the current directory."
+    datasets = get_available_datasets()
+    if datasets:
+        return "**Available CSV files:**\n\n" + "\n".join([f"- {file}" for file in datasets])
+    else:
+        return "No CSV files found in the current directory."
+
+# Initialize tokenizer on startup
+if CURRENT_TOKENIZER is None:
+    try:
+        CURRENT_TOKENIZER = AutoTokenizer.from_pretrained("bert-base-uncased")
+        print("✅ Tokenizer initialized successfully")
+    except Exception as e:
+        print(f"⚠️ Warning: Could not initialize tokenizer: {e}")
+
+print("🚀 Launching BERT Complaint Classifier...")
+print("📍 Available at: http://localhost:7860")
+
+# The entire Gradio UI definition must be within a single block
+with gr.Blocks(title="BERT Complaint Classifier", theme=gr.themes.Soft()) as app:
+    gr.Markdown("# BERT Complaint Classifier 🗣️🤖")
+    gr.Markdown("Fine-tune a BERT model or use an existing one to classify customer complaints.")
+
+    with gr.Tab("Fine-tune Model"):
+        gr.Markdown("## 🏋️ Fine-tune a New Model")
+        
+        with gr.Column(variant="panel"):
+            gr.Markdown("### 🛠️ Training Configuration")
+            with gr.Row():
+                uploaded_file = gr.File(label="Upload Training CSV File", type="filepath", file_types=["csv"])
+                preview_btn = gr.Button("Preview Dataset")
+                preview_output = gr.Markdown("Dataset info will appear here")
             
-            # Initialize the display
-            refresh_datasets_btn.click(
-                display_available_datasets,
-                outputs=available_datasets
-            )
+            with gr.Row():
+                text_column_input = gr.Textbox(label="Text Column Name", value="complaint")
+                label_column_input = gr.Textbox(label="Label Column Name", value="category")
+            
+            gr.Markdown("---")
+            
+            with gr.Row():
+                num_epochs_slider = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of Epochs")
+                batch_size_slider = gr.Slider(minimum=4, maximum=32, step=4, value=8, label="Batch Size")
+                learning_rate_slider = gr.Slider(minimum=1e-6, maximum=1e-4, step=1e-6, value=2e-5, label="Learning Rate")
             
-            # Show datasets on load
-            app.load(display_available_datasets, outputs=available_datasets)
+            gr.Markdown("---")
             
+            gr.Markdown("### ☁️ Hugging Face Hub (Optional)")
+            with gr.Row():
+                push_to_hub_checkbox = gr.Checkbox(label="Push to Hugging Face Hub")
+                hf_token_input = gr.Textbox(label="Hugging Face Token", type="password")
+            with gr.Row():
+                hf_username_input = gr.Textbox(label="Hugging Face Username")
+                hf_model_name_input = gr.Textbox(label="Model Name (for Hub)", value="bert-complaint-classifier")
+
+            train_btn = gr.Button("🚀 Start Training", variant="primary")
+            
+        gr.Markdown("---")
+        
+        training_log_output = gr.Textbox(label="Training Logs", lines=20, max_lines=20, interactive=False)
+        
+    with gr.Tab("Predict"):
+        gr.Markdown("## 🔮 Make Predictions")
+        gr.Markdown("Choose a method to classify complaints.")
+
+        with gr.Tab("Predict Single Text"):
+            with gr.Column(variant="panel"):
+                gr.Markdown("### Classify a Single Complaint")
+                
+                model_path_input = gr.Textbox(
+                    label="Model Path or Hub ID", 
+                    value="bert-base-uncased", 
+                    placeholder="e.g., local-model or your_username/your_model"
+                )
+                
+                with gr.Row():
+                    text_input = gr.Textbox(label="Complaint Text", lines=3)
+                    token_count_output = gr.Markdown("Token count: 0/512")
+                
+                predict_btn = gr.Button("Classify Complaint", variant="primary")
+                single_prediction_output = gr.Markdown("Prediction will appear here...")
+                
+                # Link token count to text input
+                text_input.change(count_tokens, inputs=text_input, outputs=token_count_output)
+            
+        with gr.Tab("Predict from CSV"):
+            with gr.Column(variant="panel"):
+                gr.Markdown("### Classify Complaints from a CSV File")
+                csv_file_input = gr.File(label="Upload CSV File (with 'complaint' column)")
+                csv_model_path = gr.Textbox(
+                    label="Model Path or Hub ID", 
+                    value="local-model",
+                    placeholder="e.g., local-model or your_username/your_model"
+                )
+                csv_predict_btn = gr.Button("Run Predictions on CSV", variant="primary")
+                csv_prediction_output = gr.Markdown("Predictions will appear here...")
+                download_link = gr.File(label="Download Full Predictions", interactive=False)
+                
+        # Link prediction buttons to functions
+        predict_btn.click(
+            predict_text, 
+            inputs=[text_input, model_path_input], 
+            outputs=single_prediction_output
+        )
+        csv_predict_btn.click(
+            predict_csv, 
+            inputs=[csv_file_input, csv_model_path], 
+            outputs=[csv_prediction_output, download_link]
+        )
+        
+    with gr.Tab("Tools"):
+        gr.Markdown("## 🔧 Tools")
+        gr.Markdown("Utilities for managing datasets and models.")
+
+        with gr.Accordion("Dataset Information"):
+            available_datasets = gr.Markdown("No CSV files found in the current directory.")
+            refresh_datasets_btn = gr.Button("🔄 Refresh Available Datasets")
             gr.Markdown("### Dataset Format Requirements")
             gr.Markdown("""
             **For training, your CSV file should have:**
@@ -676,7 +777,6 @@ def display_available_datasets():
             "Poor TV signal",TV-Radio
             ```
             """)
-            
             gr.Markdown("### Model Categories")
             categories_info = f"""
             **The model classifies complaints into these categories:**
@@ -689,23 +789,60 @@ def display_available_datasets():
             """
             gr.Markdown(categories_info)
 
-# Launch the app
-if __name__ == "__main__":
-    # Initialize tokenizer on startup
-    if CURRENT_TOKENIZER is None:
-        try:
-            CURRENT_TOKENIZER = AutoTokenizer.from_pretrained("bert-base-uncased")
-            print("✅ Tokenizer initialized successfully")
-        except Exception as e:
-            print(f"⚠️ Warning: Could not initialize tokenizer: {e}")
+        with gr.Accordion("Push Local Model to Hub"):
+            gr.Markdown("Use this to manually push a locally trained model (`./local-model`) to the Hub.")
+            with gr.Row():
+                hub_username_input_push = gr.Textbox(label="Hugging Face Username")
+                hub_model_name_input_push = gr.Textbox(label="Model Name")
+                hub_token_input_push = gr.Textbox(label="Hugging Face Token", type="password")
+            
+            push_btn = gr.Button("🚀 Push Model to Hub", variant="primary")
+            push_output = gr.verse("Results will appear here...")
+            
+            # Link the push button
+            push_btn.click(
+                push_to_hub_after_training,
+                inputs=[gr.Textbox(value=MODEL_PATH, visible=False), hub_username_input_push, hub_model_name_input_push, hub_token_input_push],
+                outputs=push_output
+            )
+
+    # All button clicks and UI logic now correctly indented within the app block
+    preview_btn.click(
+        preview_dataset, 
+        inputs=[uploaded_file, text_column_input, label_column_input], 
+        outputs=preview_output
+    )
     
-    print("🚀 Launching BERT Complaint Classifier...")
-    print("📍 Available at: http://localhost:7860")
+    train_btn.click(
+        train_model_inline,
+        inputs=[
+            uploaded_file,
+            text_column_input,
+            label_column_input,
+            num_epochs_slider,
+            batch_size_slider,
+            learning_rate_slider,
+            hf_token_input,
+            push_to_hub_checkbox,
+            hf_username_input,
+            hf_model_name_input,
+        ],
+        outputs=training_log_output,
+    )
+    
+    refresh_datasets_btn.click(
+        display_available_datasets,
+        outputs=available_datasets
+    )
     
+    # Show datasets on load
+    app.load(display_available_datasets, outputs=available_datasets)
+
+# Launch the app
+if __name__ == "__main__":
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,
         share=False,
         show_error=True
-    )
-
+    )