Update app.py

app.py

@@ -19,7 +19,6 @@ class MultiClientThemeClassifier:
         self.client_themes = {}
         self.model_loaded = False
         self.default_model = 'Qwen/Qwen3-Embedding-0.6B'
-        # CORRECTED: Add attribute to remember the last loaded model's name
         self.current_model_name = self.default_model
         
     def load_model(self, model_name: str):
@@ -36,7 +35,6 @@ class MultiClientThemeClassifier:
             print(f"Loading model: {model_name} onto CUDA device")
             self.model = SentenceTransformer(model_name, device='cuda', trust_remote_code=True)
             self.model_loaded = True
-            # CORRECTED: Remember the name of the successfully loaded model
             self.current_model_name = model_name
             return f"✅ Model '{model_name}' loaded successfully onto GPU!"
         except Exception as e:
@@ -48,7 +46,6 @@ class MultiClientThemeClassifier:
         """Internal helper to load the correct model if it's not already loaded."""
         if not self.model_loaded:
             print(f"Model not loaded. Automatically loading last selected model: {self.current_model_name}...")
-            # CORRECTED: Load the last selected model, not the default one
             status = self.load_model(self.current_model_name)
             if "Error" in status:
                 return status
@@ -93,11 +90,16 @@ class MultiClientThemeClassifier:
         except Exception as e:
             return f"Error: {str(e)}", 0.0, {}
     
-    def benchmark_csv(self, csv_filepath: str, client_id: str) -> Tuple[str, Optional[str], Optional[str]]:
-        """Benchmark the model on a CSV file, trying multiple encodings."""
-        error_status = self._ensure_model_is_loaded()
-        if error_status: return f"❌ Model could not be loaded: {error_status}", None, None
+    # CORRECTED: The benchmark function now takes the model_name as an argument
+    def benchmark_csv(self, csv_filepath: str, client_id: str, model_name: str) -> Tuple[str, Optional[str], Optional[str]]:
+        """Benchmark a specific model on a CSV file."""
+        # Step 1: Explicitly load the model requested by the user for this benchmark run.
+        load_status = self.load_model(model_name)
+        # We allow the function to proceed if the model is "already loaded", but stop for any other error.
+        if "❌" in load_status:
+             return f"❌ Model '{model_name}' could not be loaded for benchmarking.\n\nError: {load_status}", None, None
         
+        # Step 2: Proceed with the benchmark logic as before.
         encodings_to_try = ['utf-8-sig', 'utf-8', 'cp1256', 'latin1']
         df = None
         for encoding in encodings_to_try:
@@ -170,13 +172,17 @@ def classify_interface(text: str, client_id: str, confidence_threshold: float):
     
     return result, ""
 
+# CORRECTED: The interface now accepts model_name
 @spaces.GPU(duration=300)
-def benchmark_interface(csv_file_obj, client_id: str):
+def benchmark_interface(csv_file_obj, client_id: str, model_name: str):
     if csv_file_obj is None:
         return "Please upload a CSV file!", None, None
+    if not model_name.strip():
+        return "Please enter a model name for the benchmark!", None, None
     try:
         csv_filepath = csv_file_obj.name
-        return classifier.benchmark_csv(csv_filepath, client_id)
+        # Pass the model name from the UI down to the classifier method
+        return classifier.benchmark_csv(csv_filepath, client_id, model_name.strip())
     except Exception as e:
         error_details = traceback.format_exc()
         return f"❌ Error processing CSV file object: {str(e)}\n\nDetails:\n{error_details}", None, None
@@ -187,8 +193,9 @@ with gr.Blocks(title="Custom Themes Classification MVP", theme=gr.themes.Soft())
     
     with gr.Tab("🚀 Setup & Model"):
         gr.Markdown("### Step 1: Load the Embedding Model (Optional)")
-        gr.Markdown("A default model (`Qwen/Qwen3-Embedding-0.6B`) will load automatically on first use.")
+        gr.Markdown("A default model (`Qwen/Qwen3-Embedding-0.6B`) will load automatically on first use. You can specify a different model here to use it in other tabs.")
         with gr.Row():
+            # This input is now used by the benchmark tab as well
             model_input = gr.Textbox(label="HuggingFace Model Name", value="Qwen/Qwen3-Embedding-0.6B")
             load_btn = gr.Button("Load Model", variant="primary")
         load_status = gr.Textbox(label="Status", interactive=False)
@@ -215,7 +222,7 @@ with gr.Blocks(title="Custom Themes Classification MVP", theme=gr.themes.Soft())
         classify_btn.click(classify_interface, inputs=[text_input, client_select, confidence_slider], outputs=[classification_result, gr.Textbox(visible=False)])
     
     with gr.Tab("📊 CSV Benchmarking"):
-        gr.Markdown("### Benchmark on Your Dataset\nUpload a CSV with `text` and `real_tag` columns.")
+        gr.Markdown("### Benchmark on Your Dataset\nUpload a CSV with `text` and `real_tag` columns. The model from the 'Setup & Model' tab will be loaded and used for the benchmark.")
         with gr.Row():
             with gr.Column():
                 csv_upload = gr.File(label="Upload CSV File", file_types=[".csv"])
@@ -226,7 +233,13 @@ with gr.Blocks(title="Custom Themes Classification MVP", theme=gr.themes.Soft())
         with gr.Row():
             results_csv = gr.File(label="Download Detailed Results", interactive=False)
             visualization = gr.HTML(label="Visualization")
-        benchmark_btn.click(benchmark_interface, inputs=[csv_upload, benchmark_client], outputs=[benchmark_results, results_csv, visualization])
+        
+        # CORRECTED: The button now sends the model_input value to the benchmark function
+        benchmark_btn.click(
+            benchmark_interface, 
+            inputs=[csv_upload, benchmark_client, model_input], 
+            outputs=[benchmark_results, results_csv, visualization]
+        )
 
 # Launch the app
 if __name__ == "__main__":