Update app.py

app.py

@@ -3,183 +3,423 @@ import pandas as pd
 import torch
 from sentence_transformers import SentenceTransformer, util
 import numpy as np
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 import io
 import plotly.express as px
+import plotly.graph_objects as go
+from collections import defaultdict
+import json
 import traceback
-import spaces
-import os
-
-# THE FIX IS HERE: Use a relative path for the cache directory.
-# This creates a writable 'persistent_cache' folder inside the app's own directory.
-os.environ['SENTENCE_TRANSFORMERS_HOME'] = './persistent_cache'
 
 class MultiClientThemeClassifier:
     def __init__(self):
-        self.model: SentenceTransformer | None = None
-        self.model_name: str | None = None
-        self.client_themes = {}
-
-    def _ensure_model_loaded(self):
-        """
-        Checks if the model is loaded in the current process.
-        If not, it reloads it using the saved model_name.
-        """
-        if self.model is None:
-            if self.model_name is None:
-                raise ValueError("Model name not set. Please go to the 'Setup & Model' tab and load a model first.")
-            
-            print(f"Model not found in current process. Reloading '{self.model_name}' from cache...")
-            try:
-                self.model = SentenceTransformer(self.model_name)
-                print(f"Model '{self.model_name}' reloaded successfully.")
-            except Exception as e:
-                print(f"FATAL: Failed to reload model '{self.model_name}': {e}")
-                raise e
-
-    @spaces.GPU
-    def load_model(self, model_name: str):
-        """Loads the model and saves its name to the state."""
+        self.model = None
+        self.client_themes = {}  # {client_id: {theme: prototype_embedding}}
+        self.model_loaded = False
+        
+    def load_model(self, model_name: str = 'Qwen/Qwen3-Embedding-0.6B'):
+        """Load the embedding model"""
         try:
+            if self.model_loaded:
+                # If switching models, reset everything
+                self.model = None
+                self.client_themes = {}
+                self.model_loaded = False
+            
             print(f"Loading model: {model_name}")
-            self.model = SentenceTransformer(model_name)
-            self.model_name = model_name
-            self.client_themes = {}
+            self.model = SentenceTransformer(model_name,trust_remote_code=True)
+            self.model_loaded = True
             return f"✅ Model '{model_name}' loaded successfully!"
         except Exception as e:
-            self.model = None
-            self.model_name = None
-            return f"❌ Error loading model '{model_name}': {traceback.format_exc()}"
-
-    @spaces.GPU
-    def add_client_themes(self, client_id: str, themes: List[str]):
-        """Adds themes for a client, ensuring the model is loaded first."""
+            self.model_loaded = False
+            error_details = traceback.format_exc()
+            return f"❌ Error loading model '{model_name}': {str(e)}\n\nDetails:\n{error_details}"
+    
+    def add_client_themes(self, client_id: str, themes: List[str], examples_per_theme: Dict[str, List[str]] = None):
+        """Add themes for a specific client"""
+        if not self.model_loaded:
+            return "❌ Please load the model first!"
+        
         try:
-            self._ensure_model_loaded()
             self.client_themes[client_id] = {}
-            prototypes = self.model.encode(themes, convert_to_tensor=True)
-            for theme, prototype in zip(themes, prototypes):
+            
+            for theme in themes:
+                if examples_per_theme and theme in examples_per_theme:
+                    # Use provided examples to create prototype
+                    examples = examples_per_theme[theme]
+                    embeddings = self.model.encode(examples, convert_to_tensor=True)
+                    prototype = torch.mean(embeddings, dim=0)
+                else:
+                    # Use theme name itself as prototype (fallback)
+                    prototype = self.model.encode(theme, convert_to_tensor=True)
+                
                 self.client_themes[client_id][theme] = prototype
+            
             return f"✅ Added {len(themes)} themes for client '{client_id}'"
         except Exception as e:
             return f"❌ Error adding themes: {str(e)}"
-
-    @spaces.GPU
-    def benchmark_csv(self, csv_content: str, client_id: str) -> Tuple[str, pd.DataFrame | None, str | None]:
-        """Benchmarks a CSV, ensuring the model is loaded first."""
+    
+    def classify_text(self, text: str, client_id: str, confidence_threshold: float = 0.3) -> Tuple[str, float, Dict[str, float]]:
+        """Classify a single text for a specific client"""
+        if not self.model_loaded:
+            return "Model not loaded", 0.0, {}
+        
+        if client_id not in self.client_themes:
+            return "Client not found", 0.0, {}
+        
+        try:
+            # Encode input text
+            text_embedding = self.model.encode(text, convert_to_tensor=True)
+            
+            # Calculate similarities with all themes
+            similarities = {}
+            for theme, prototype in self.client_themes[client_id].items():
+                similarity = util.cos_sim(text_embedding, prototype).item()
+                similarities[theme] = similarity
+            
+            # Get best match
+            best_theme = max(similarities, key=similarities.get)
+            best_score = similarities[best_theme]
+            
+            # Apply confidence threshold
+            if best_score < confidence_threshold:
+                return "UNKNOWN_THEME", best_score, similarities
+            
+            return best_theme, best_score, similarities
+        except Exception as e:
+            return f"Error: {str(e)}", 0.0, {}
+    
+    def benchmark_csv(self, csv_content: str, client_id: str) -> Tuple[str, str, str]:
+        """Benchmark the model on a CSV file"""
+        if not self.model_loaded:
+            return "❌ Model not loaded!", "", ""
+        
         try:
-            self._ensure_model_loaded()
-            print("Model confirmed loaded in benchmark process. Starting benchmark...")
+            print("Starting CSV benchmark...")
             
+            # Read CSV
             df = pd.read_csv(io.StringIO(csv_content))
+            print(f"CSV loaded with shape: {df.shape}")
+            print(f"CSV columns: {df.columns.tolist()}")
+            
+            # Validate CSV format
             if 'text' not in df.columns or 'real_tag' not in df.columns:
-                return "❌ CSV must have 'text' and 'real_tag' columns!", None, ""
-
-            df.dropna(subset=['text', 'real_tag'], inplace=True)
+                return "❌ CSV must have 'text' and 'real_tag' columns!", "", ""
             
+            # Clean data
+            df = df.dropna(subset=['text', 'real_tag'])
+            df['text'] = df['text'].astype(str)
+            df['real_tag'] = df['real_tag'].astype(str)
+            print(f"After cleaning: {df.shape}")
+            
+            # Get unique themes from CSV
             unique_themes = df['real_tag'].unique().tolist()
-            self.add_client_themes(client_id, unique_themes)
+            print(f"Unique themes found: {len(unique_themes)} - {unique_themes}")
             
-            texts_to_classify = df['text'].astype(str).tolist()
-            text_embeddings = self.model.encode(texts_to_classify, convert_to_tensor=True, show_progress_bar=True)
+            # Add themes for this client (using theme names as prototypes for demo)
+            theme_add_result = self.add_client_themes(client_id, unique_themes)
+            print(f"Theme addition result: {theme_add_result}")
             
-            themes = list(self.client_themes[client_id].keys())
-            prototypes = torch.stack(list(self.client_themes[client_id].values()))
+            # Classify all texts with progress
+            predictions = []
+            confidences = []
             
-            similarities_matrix = util.cos_sim(text_embeddings, prototypes)
+            print("Starting classification...")
+            for idx, row in df.iterrows():
+                try:
+                    text = str(row['text'])[:500]  # Limit text length
+                    pred_theme, confidence, _ = self.classify_text(text, client_id)
+                    predictions.append(pred_theme)
+                    confidences.append(confidence)
+                    
+                    if idx % 10 == 0:  # Progress logging
+                        print(f"Processed {idx + 1}/{len(df)} samples")
+                        
+                except Exception as e:
+                    print(f"Error classifying row {idx}: {str(e)}")
+                    predictions.append("ERROR")
+                    confidences.append(0.0)
             
-            best_scores, best_indices = torch.max(similarities_matrix, dim=1)
+            print("Classification complete!")
             
-            df['predicted_tag'] = [themes[i] for i in best_indices]
-            df['confidence'] = best_scores.tolist()
-
-            correct = (df['real_tag'] == df['predicted_tag']).sum()
-            total = len(df)
+            df['predicted_tag'] = predictions
+            df['confidence'] = confidences
+            
+            # Calculate metrics
+            valid_predictions = df[df['predicted_tag'] != 'ERROR']
+            correct = (valid_predictions['real_tag'] == valid_predictions['predicted_tag']).sum()
+            total = len(valid_predictions)
             accuracy = correct / total if total > 0 else 0
+            
+            print(f"Metrics calculated: {correct}/{total} = {accuracy:.2%}")
+            
+            # Generate results summary
+            results_summary = f"""
+📊 **Benchmarking Results**
 
-            results_summary = f"📊 **Benchmarking Results**\n\n- **Accuracy: {accuracy:.2%}** ({correct} / {total} correct)"
-            fig = px.bar(df['real_tag'].value_counts(), title="Theme Distribution in Dataset")
-            visualization_html = fig.to_html()
-
-            return results_summary, df, visualization_html
+**Overall Metrics:**
+- Total samples: {total}
+- Correct predictions: {correct}
+- **Accuracy: {accuracy:.2%}**
+- Average confidence: {np.mean([c for c in confidences if c > 0]):.3f}
 
-        except ValueError as e:
-            return f"❌ {str(e)}", None, ""
+**Per-Theme Breakdown:**
+"""
+            
+            for theme in unique_themes:
+                theme_df = valid_predictions[valid_predictions['real_tag'] == theme]
+                if len(theme_df) > 0:
+                    theme_correct = (theme_df['real_tag'] == theme_df['predicted_tag']).sum()
+                    theme_total = len(theme_df)
+                    theme_acc = theme_correct / theme_total if theme_total > 0 else 0
+                    avg_conf = theme_df['confidence'].mean()
+                    
+                    results_summary += f"- **{theme}**: {theme_acc:.2%} ({theme_correct}/{theme_total}) - Avg conf: {avg_conf:.3f}\n"
+            
+            # Create simple visualization
+            try:
+                theme_counts = [len(df[df['real_tag'] == theme]) for theme in unique_themes]
+                fig = px.bar(
+                    x=unique_themes,
+                    y=theme_counts,
+                    title="Theme Distribution in Dataset",
+                    labels={'x': 'Themes', 'y': 'Count'}
+                )
+                visualization_html = fig.to_html()
+            except Exception as viz_error:
+                print(f"Visualization error: {viz_error}")
+                visualization_html = "<p>Visualization error occurred</p>"
+            
+            # Save CSV to a temporary file for download
+            import tempfile
+            import os
+            
+            temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8')
+            df.to_csv(temp_file.name, index=False)
+            temp_file.close()
+            
+            return results_summary, temp_file.name, visualization_html
+            
         except Exception as e:
-            return f"❌ Error during benchmarking: {traceback.format_exc()}", None, ""
-
+            error_details = traceback.format_exc()
+            print(f"Full error: {error_details}")
+            return f"❌ Error during benchmarking: {str(e)}\n\nFull traceback:\n{error_details}", "", ""
 
-# --- Interface Functions ---
+# Initialize the classifier
+classifier = MultiClientThemeClassifier()
 
-def load_model_interface(classifier, model_name: str):
+def load_model_interface(model_name: str):
     if not model_name.strip():
-        # Fallback to a default model if input is empty
-        model_name = 'sentence-transformers/all-MiniLM-L6-v2'
-    status = classifier.load_model(model_name.strip())
-    return status, classifier
+        model_name = 'Qwen/Qwen3-Embedding-0.6B'  # Default
+    return classifier.load_model(model_name.strip())
 
-def add_themes_interface(classifier, client_id: str, themes_text: str):
-    if not client_id.strip() or not themes_text.strip():
-        return "❌ Client ID and Themes cannot be empty.", classifier
+def add_themes_interface(client_id: str, themes_text: str):
+    if not themes_text.strip():
+        return "❌ Please enter themes!"
+    
     themes = [theme.strip() for theme in themes_text.split('\n') if theme.strip()]
-    status = classifier.add_client_themes(client_id, themes)
-    return status, classifier
+    return classifier.add_client_themes(client_id, themes)
+
+def classify_interface(text: str, client_id: str, confidence_threshold: float):
+    if not text.strip():
+        return "Please enter text to classify!", ""
+    
+    pred_theme, confidence, similarities = classifier.classify_text(text, client_id, confidence_threshold)
+    
+    # Format similarities for display
+    sim_display = "**Similarity Scores:**\n"
+    sorted_sims = sorted(similarities.items(), key=lambda x: x[1], reverse=True)
+    for theme, sim in sorted_sims:
+        sim_display += f"- {theme}: {sim:.3f}\n"
+    
+    result = f"""
+🎯 **Predicted Theme:** {pred_theme}
+🔥 **Confidence:** {confidence:.3f}
+
+{sim_display}
+"""
+    
+    return result, ""
 
-def benchmark_interface(classifier, csv_file, client_id: str):
+def benchmark_interface(csv_file, client_id: str):
     if csv_file is None:
-        return "Please upload a CSV file!", None, "", classifier
-    if not client_id.strip():
-        return "❌ Please enter a Client ID for the benchmark.", None, "", classifier
+        return "Please upload a CSV file!", "", ""
+    
     try:
-        with open(csv_file.name, 'r', encoding='utf-8') as f:
-            csv_content = f.read()
-        results, df, viz = classifier.benchmark_csv(csv_content, client_id)
-        return results, df, viz, classifier
+        # Handle both file objects and file paths
+        if hasattr(csv_file, 'read'):
+            csv_content = csv_file.read().decode('utf-8')
+        elif hasattr(csv_file, 'name'):
+            with open(csv_file.name, 'r', encoding='utf-8') as f:
+                csv_content = f.read()
+        else:
+            csv_content = str(csv_file)
+            
+        return classifier.benchmark_csv(csv_content, client_id)
     except Exception as e:
-        return f"❌ Error processing CSV: {traceback.format_exc()}", None, "", classifier
+        error_details = traceback.format_exc()
+        return f"❌ Error reading CSV: {str(e)}\n\nDetails:\n{error_details}", "", ""
 
-# --- Gradio UI ---
+# Create the Gradio interface
 with gr.Blocks(title="Custom Themes Classification MVP", theme=gr.themes.Soft()) as demo:
-    classifier_state = gr.State(MultiClientThemeClassifier())
+    gr.Markdown("""
+    # 🎯 Custom Themes Classification - MVP
     
-    gr.Markdown("# 🎯 Custom Themes Classification - MVP")
+    **A scalable, cost-effective solution for multi-client theme classification**
+    
+    This demo showcases an embedding-based approach that can:
+    - ✅ Handle multiple clients with different themes
+    - ✅ Distinguish between similar themes (e.g., "Real Estate Financing" vs "Personal Financing")  
+    - ✅ Process ~1M posts/day at low cost (~$500/month vs $30k/month for pure LLM)
+    - ✅ Provide confidence scores and similarity breakdowns
+    """)
     
     with gr.Tab("🚀 Setup & Model"):
-        model_input = gr.Textbox(label="HuggingFace Model Name", value="Qwen/Qwen3-Embedding-0.6B")
-        load_btn = gr.Button("Load Model", variant="primary")
+        gr.Markdown("### Step 1: Load the Embedding Model")
+        
+        with gr.Row():
+            model_input = gr.Textbox(
+                label="HuggingFace Model Name",
+                value="Qwen/Qwen3-Embedding-0.6B",
+                placeholder="e.g., sentence-transformers/all-MiniLM-L6-v2",
+                info="Enter any SentenceTransformer-compatible model from HuggingFace"
+            )
+            load_btn = gr.Button("Load Model", variant="primary")
+        
         load_status = gr.Textbox(label="Status", interactive=False)
         
-        gr.Markdown("---")
+        gr.Markdown("""
+        **Popular Models:**
+        - `Qwen/Qwen3-Embedding-0.6B` - High quality, multilingual
+        - `sentence-transformers/all-MiniLM-L6-v2` - Fast, lightweight  
+        - `sentence-transformers/all-mpnet-base-v2` - High accuracy
+        - `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2` - Multilingual
+        - `intfloat/multilingual-e5-base` - Strong multilingual performance
+        """)
+        
+        load_btn.click(load_model_interface, inputs=[model_input], outputs=load_status)
+        
+        gr.Markdown("### Step 2: Add Themes for a Client")
+        with gr.Row():
+            client_input = gr.Textbox(label="Client ID", placeholder="e.g., client_1")
+            themes_input = gr.Textbox(
+                label="Themes (one per line)", 
+                lines=5,
+                placeholder="e.g.:\nReal Estate Financing\nPersonal Financing\nPrivate Education\nSports"
+            )
         
-        client_input = gr.Textbox(label="Client ID", placeholder="e.g., client_1")
-        themes_input = gr.Textbox(label="Themes (one per line)", lines=5)
         add_themes_btn = gr.Button("Add Themes", variant="secondary")
         themes_status = gr.Textbox(label="Status", interactive=False)
-
+        
+        add_themes_btn.click(
+            add_themes_interface,
+            inputs=[client_input, themes_input],
+            outputs=themes_status
+        )
+    
+    with gr.Tab("🔍 Single Text Classification"):
+        gr.Markdown("### Classify Individual Posts")
+        
+        with gr.Row():
+            with gr.Column():
+                text_input = gr.Textbox(
+                    label="Text to Classify",
+                    lines=3,
+                    placeholder="Enter text to classify..."
+                )
+                client_select = gr.Textbox(
+                    label="Client ID",
+                    placeholder="e.g., client_1"
+                )
+                confidence_slider = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.3,
+                    step=0.1,
+                    label="Confidence Threshold"
+                )
+                classify_btn = gr.Button("Classify", variant="primary")
+            
+            with gr.Column():
+                classification_result = gr.Markdown(label="Results")
+                
+        classify_btn.click(
+            classify_interface,
+            inputs=[text_input, client_select, confidence_slider],
+            outputs=[classification_result, gr.Textbox(visible=False)]
+        )
+    
     with gr.Tab("📊 CSV Benchmarking"):
-        csv_upload = gr.File(label="Upload CSV File", file_types=[".csv"])
-        benchmark_client = gr.Textbox(label="Client ID for Benchmark", placeholder="e.g., benchmark_client")
-        benchmark_btn = gr.Button("Run Benchmark", variant="primary")
-        benchmark_results = gr.Markdown(label="Benchmark Results")
-        results_dataframe = gr.Dataframe(label="Detailed Results", interactive=False, wrap=True)
-        visualization = gr.HTML(label="Visualization")
+        gr.Markdown("""
+        ### Benchmark on Your Dataset
+        
+        Upload a CSV file with columns:
+        - `text`: The posts/content to classify
+        - `real_tag`: The correct theme labels
+        
+        The system will automatically extract unique themes and evaluate performance.
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                csv_upload = gr.File(
+                    label="Upload CSV File",
+                    file_types=[".csv"]
+                )
+                benchmark_client = gr.Textbox(
+                    label="Client ID for Benchmark",
+                    placeholder="e.g., benchmark_client"
+                )
+                benchmark_btn = gr.Button("Run Benchmark", variant="primary")
+            
+            with gr.Column():
+                benchmark_results = gr.Markdown(label="Benchmark Results")
+        
+        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]
+        )
     
-    # --- Event Handlers ---
-    load_btn.click(
-        load_model_interface,
-        inputs=[classifier_state, model_input],
-        outputs=[load_status, classifier_state]
-    )
-    add_themes_btn.click(
-        add_themes_interface,
-        inputs=[classifier_state, client_input, themes_input],
-        outputs=[themes_status, classifier_state]
-    )
-    benchmark_btn.click(
-        benchmark_interface,
-        inputs=[classifier_state, csv_upload, benchmark_client],
-        outputs=[benchmark_results, results_dataframe, visualization, classifier_state]
-    )
+    with gr.Tab("📋 About & Usage"):
+        gr.Markdown("""
+        ## 🎯 Solution Overview
+        
+        This MVP demonstrates a **hybrid embedding-based approach** for Custom Themes classification:
+        
+        ### ✅ Key Advantages:
+        1. **Cost Effective**: ~$500/month vs $30,000/month for pure LLM approach
+        2. **Fast**: Can handle 1M+ posts/day with sub-second response times  
+        3. **Multi-Client**: Each client can have completely different themes
+        4. **Disambiguates Similar Themes**: Uses semantic embeddings to distinguish between similar concepts
+        5. **Confidence Scoring**: Provides transparency in predictions
+        
+        ### 🏗️ Architecture:
+        1. **Embedding Model**: Customizable SentenceTransformer models from HuggingFace
+        2. **Theme Prototypes**: Each client's themes represented as embedding vectors  
+        3. **Similarity Matching**: Cosine similarity for classification
+        4. **Confidence Thresholding**: Flags uncertain predictions
+        
+        ### 📈 Scaling Strategy:
+        - **Batch Processing**: Process thousands of posts simultaneously
+        - **GPU Optimization**: Single GPU can handle 1M posts/day
+        - **Caching**: Store client prototypes in memory/Redis
+        - **Hybrid Fallback**: LLM backup for ambiguous cases (5-10% of posts)
+        
+        ### 🔧 Usage Instructions:
+        1. **Setup Tab**: Load model and define client themes
+        2. **Single Classification**: Test individual posts  
+        3. **CSV Benchmark**: Evaluate on your datasets
+        
+        ---
+        
+        **Scalable Theme Classification MVP**
+        """)
 
+# Launch the app
 if __name__ == "__main__":
     demo.launch(share=True)