Add remove labels functionality with admin panel UI

app.py

@@ -155,6 +155,81 @@ def get_current_stats():
     except Exception as e:
         return f"Error getting stats: {str(e)}"
 
+def get_labels_table():
+    """
+    Get all current labels as a formatted table for display.
+    """
+    if handler is None:
+        return "Handler not initialized"
+    
+    if not hasattr(handler, 'class_ids') or len(handler.class_ids) == 0:
+        return "No labels currently loaded"
+    
+    try:
+        # Create a formatted table of labels
+        table_data = []
+        for id, name in zip(handler.class_ids, handler.class_names):
+            table_data.append([int(id), name])
+        
+        return table_data
+    except Exception as e:
+        return f"Error getting labels: {str(e)}"
+
+def remove_labels_admin(admin_token, ids_to_remove_str):
+    """
+    Admin function to remove labels by ID.
+    """
+    if handler is None:
+        return "Error: Handler not initialized"
+    
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        # Parse the IDs from comma-separated string
+        if not ids_to_remove_str or ids_to_remove_str.strip() == "":
+            return "❌ Error: Please provide IDs to remove (comma-separated)"
+        
+        ids_to_remove = []
+        for id_str in ids_to_remove_str.split(','):
+            id_str = id_str.strip()
+            if id_str:
+                ids_to_remove.append(int(id_str))
+        
+        if not ids_to_remove:
+            return "❌ Error: No valid IDs provided"
+        
+        # Get names of items to be removed for confirmation
+        removed_names = []
+        if hasattr(handler, 'class_ids'):
+            for id in ids_to_remove:
+                if id in handler.class_ids:
+                    idx = handler.class_ids.index(id)
+                    removed_names.append(f"{id}: {handler.class_names[idx]}")
+        
+        result = handler({
+            "inputs": {
+                "op": "remove_labels",
+                "token": admin_token,
+                "ids": ids_to_remove
+            }
+        })
+        
+        if result.get("status") == "ok":
+            removed_list = "\n".join(removed_names) if removed_names else "None found"
+            return f"✅ Success! Removed {result.get('removed', 0)} labels. Current version: {result.get('labels_version', 'unknown')}\n\nRemoved items:\n{removed_list}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        elif result.get("error") == "no_ids_provided":
+            return "❌ Error: No IDs provided"
+        else:
+            return f"❌ Error: {result.get('detail', result.get('error', 'Unknown error'))}"
+            
+    except ValueError:
+        return "❌ Error: Invalid ID format. Please provide comma-separated numbers (e.g., 1001,1002,1003)"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
 # Create Gradio interface
 print("Creating Gradio interface...")
 with gr.Blocks(title="MobileCLIP Image Classifier") as demo:
@@ -265,6 +340,46 @@ with gr.Blocks(title="MobileCLIP Image Classifier") as demo:
                 inputs=[admin_token_input, version_input],
                 outputs=reload_output
             )
+        
+        with gr.Accordion("🗑️ Remove Labels", open=False):
+            gr.Markdown("Remove specific labels by their IDs")
+            
+            # Display current labels
+            labels_table = gr.Dataframe(
+                value=get_labels_table(),
+                headers=["ID", "Name"],
+                label="Current Labels",
+                interactive=False,
+                height=300
+            )
+            
+            refresh_labels_btn = gr.Button("🔄 Refresh Label List", size="sm")
+            refresh_labels_btn.click(
+                fn=get_labels_table,
+                inputs=[],
+                outputs=labels_table
+            )
+            
+            gr.Markdown("Enter IDs to remove (comma-separated):")
+            ids_to_remove_input = gr.Textbox(
+                label="IDs to Remove",
+                placeholder="e.g., 1001, 1002, 1003",
+                lines=1
+            )
+            
+            remove_btn = gr.Button("🗑️ Remove Selected Labels", variant="stop")
+            remove_output = gr.Markdown()
+            
+            def remove_and_refresh(token, ids):
+                result = remove_labels_admin(token, ids)
+                updated_table = get_labels_table()
+                return result, updated_table
+            
+            remove_btn.click(
+                fn=remove_and_refresh,
+                inputs=[admin_token_input, ids_to_remove_input],
+                outputs=[remove_output, labels_table]
+            )
     
     with gr.Tab("ℹ️ About"):
         gr.Markdown("""

app_backup.py

@@ -0,0 +1,280 @@
+import gradio as gr
+import base64
+import json
+import os
+from PIL import Image
+import io
+from handler import EndpointHandler
+
+handler = EndpointHandler()
+
+def classify_image(image, top_k=10):
+    """
+    Main classification function for public interface.
+    """
+    if image is None:
+        return None, "Please upload an image"
+    
+    try:
+        # Convert PIL image to base64
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        img_b64 = base64.b64encode(buffered.getvalue()).decode()
+        
+        # Call handler
+        result = handler({
+            "inputs": {
+                "image": img_b64,
+                "top_k": int(top_k)
+            }
+        })
+        
+        # Format results for display
+        if isinstance(result, list):
+            # Create formatted output
+            output_text = "**Top {} Classifications:**\n\n".format(len(result))
+            
+            # Create a dictionary for the bar chart
+            chart_data = {}
+            
+            for i, item in enumerate(result, 1):
+                score_pct = item['score'] * 100
+                output_text += f"{i}. **{item['label']}** (ID: {item['id']}): {score_pct:.2f}%\n"
+                chart_data[item['label']] = item['score']
+            
+            return chart_data, output_text
+        else:
+            return None, f"Error: {result.get('error', 'Unknown error')}"
+            
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def upsert_labels_admin(admin_token, new_items_json):
+    """
+    Admin function to add new labels.
+    """
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        # Parse the JSON input
+        items = json.loads(new_items_json) if new_items_json else []
+        
+        result = handler({
+            "inputs": {
+                "op": "upsert_labels",
+                "token": admin_token,
+                "items": items
+            }
+        })
+        
+        if result.get("status") == "ok":
+            return f"✅ Success! Added {result.get('added', 0)} new labels. Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        else:
+            return f"❌ Error: {result.get('detail', result.get('error', 'Unknown error'))}"
+            
+    except json.JSONDecodeError:
+        return "❌ Error: Invalid JSON format"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def reload_labels_admin(admin_token, version):
+    """
+    Admin function to reload a specific label version.
+    """
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        result = handler({
+            "inputs": {
+                "op": "reload_labels",
+                "token": admin_token,
+                "version": int(version) if version else 1
+            }
+        })
+        
+        if result.get("status") == "ok":
+            return f"✅ Labels reloaded successfully! Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("status") == "nochange":
+            return f"ℹ️ No change needed. Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        elif result.get("error") == "invalid_version":
+            return "❌ Error: Invalid version number"
+        else:
+            return f"❌ Error: {result.get('error', 'Unknown error')}"
+            
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def get_current_stats():
+    """
+    Get current label statistics.
+    """
+    try:
+        num_labels = len(handler.class_ids) if hasattr(handler, 'class_ids') else 0
+        version = getattr(handler, 'labels_version', 1)
+        device = handler.device if hasattr(handler, 'device') else "unknown"
+        
+        stats = f"""
+        **Current Statistics:**
+        - Number of labels: {num_labels}
+        - Labels version: {version}
+        - Device: {device}
+        - Model: MobileCLIP-B
+        """
+        
+        if hasattr(handler, 'class_names') and len(handler.class_names) > 0:
+            stats += f"\n- Sample labels: {', '.join(handler.class_names[:5])}"
+            if len(handler.class_names) > 5:
+                stats += "..."
+        
+        return stats
+    except Exception as e:
+        return f"Error getting stats: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="MobileCLIP Image Classifier") as demo:
+    gr.Markdown("""
+    # 🖼️ MobileCLIP-B Zero-Shot Image Classifier
+    
+    Upload an image to classify it using MobileCLIP-B model with dynamic label management.
+    """)
+    
+    with gr.Tab("🔍 Image Classification"):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    type="pil",
+                    label="Upload Image"
+                )
+                top_k_slider = gr.Slider(
+                    minimum=1,
+                    maximum=50,
+                    value=10,
+                    step=1,
+                    label="Number of top results to show"
+                )
+                classify_btn = gr.Button("🚀 Classify Image", variant="primary")
+            
+            with gr.Column():
+                output_chart = gr.BarPlot(
+                    label="Classification Confidence",
+                    x_label="Label",
+                    y_label="Confidence",
+                    vertical=False,
+                    height=400
+                )
+                output_text = gr.Markdown(label="Classification Results")
+        
+        gr.Examples(
+            examples=[
+                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/image_classifier/examples/cheetah.jpg"],
+                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/image_classifier/examples/elephant.jpg"],
+                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/image_classifier/examples/giraffe.jpg"]
+            ],
+            inputs=input_image,
+            label="Example Images"
+        )
+        
+        classify_btn.click(
+            classify_image,
+            inputs=[input_image, top_k_slider],
+            outputs=[output_chart, output_text]
+        )
+    
+    with gr.Tab("🔧 Admin Panel"):
+        gr.Markdown("""
+        ### Admin Functions
+        **Note:** Requires admin token (set via environment variable `ADMIN_TOKEN`)
+        """)
+        
+        with gr.Row():
+            admin_token_input = gr.Textbox(
+                label="Admin Token",
+                type="password",
+                placeholder="Enter admin token"
+            )
+        
+        with gr.Accordion("📊 Current Statistics", open=True):
+            stats_display = gr.Markdown(value=get_current_stats())
+            refresh_stats_btn = gr.Button("🔄 Refresh Stats")
+            refresh_stats_btn.click(
+                get_current_stats,
+                outputs=stats_display
+            )
+        
+        with gr.Accordion("➕ Add New Labels", open=False):
+            gr.Markdown("""
+            Add new labels by providing JSON array:
+            ```json
+            [
+                {"id": 100, "name": "new_object", "prompt": "a photo of a new_object"},
+                {"id": 101, "name": "another_object", "prompt": "a photo of another_object"}
+            ]
+            ```
+            """)
+            new_items_input = gr.Code(
+                label="New Items JSON",
+                language="json",
+                lines=5,
+                value='[\n  {"id": 100, "name": "example", "prompt": "a photo of example"}\n]'
+            )
+            upsert_btn = gr.Button("➕ Add Labels", variant="primary")
+            upsert_output = gr.Markdown()
+            
+            upsert_btn.click(
+                upsert_labels_admin,
+                inputs=[admin_token_input, new_items_input],
+                outputs=upsert_output
+            )
+        
+        with gr.Accordion("🔄 Reload Label Version", open=False):
+            gr.Markdown("Reload labels from a specific version stored in the Hub")
+            version_input = gr.Number(
+                label="Version Number",
+                value=1,
+                precision=0
+            )
+            reload_btn = gr.Button("🔄 Reload Version", variant="primary")
+            reload_output = gr.Markdown()
+            
+            reload_btn.click(
+                reload_labels_admin,
+                inputs=[admin_token_input, version_input],
+                outputs=reload_output
+            )
+    
+    with gr.Tab("ℹ️ About"):
+        gr.Markdown("""
+        ## About MobileCLIP-B Classifier
+        
+        This Space provides a web interface for Apple's MobileCLIP-B model, optimized for fast zero-shot image classification.
+        
+        ### Features:
+        - 🚀 **Fast inference**: < 30ms on GPU
+        - 🏷️ **Dynamic labels**: Add/update labels without redeployment
+        - 🔄 **Version control**: Track and reload label versions
+        - 📊 **Visual results**: Bar charts and confidence scores
+        
+        ### Environment Variables (set in Space Settings):
+        - `ADMIN_TOKEN`: Secret token for admin operations
+        - `HF_LABEL_REPO`: Hub repository for label storage (e.g., "username/labels")
+        - `HF_WRITE_TOKEN`: Token with write permissions to label repo
+        - `HF_READ_TOKEN`: Token with read permissions (optional, defaults to write token)
+        
+        ### Model Details:
+        - **Architecture**: MobileCLIP-B with MobileOne blocks
+        - **Text Encoder**: Transformer-based, 77 token context
+        - **Image Size**: 224x224
+        - **Embedding Dim**: 512
+        
+        ### License:
+        Model weights are licensed under Apple Sample Code License (ASCL).
+        """)
+
+if __name__ == "__main__":
+    demo.launch()

app_fixed.py

@@ -0,0 +1,301 @@
+import gradio as gr
+import base64
+import json
+import os
+from PIL import Image
+import io
+from handler import EndpointHandler
+
+# Initialize handler
+print("Initializing MobileCLIP handler...")
+try:
+    handler = EndpointHandler()
+    print(f"Handler initialized successfully! Device: {handler.device}")
+except Exception as e:
+    print(f"Error initializing handler: {e}")
+    handler = None
+
+def classify_image(image, top_k=10):
+    """
+    Main classification function for public interface.
+    """
+    if handler is None:
+        return "Error: Handler not initialized", None
+        
+    if image is None:
+        return "Please upload an image", None
+    
+    try:
+        # Convert PIL image to base64
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        img_b64 = base64.b64encode(buffered.getvalue()).decode()
+        
+        # Call handler
+        result = handler({
+            "inputs": {
+                "image": img_b64,
+                "top_k": int(top_k)
+            }
+        })
+        
+        # Format results for display
+        if isinstance(result, list):
+            # Create formatted output
+            output_text = "**Top {} Classifications:**\n\n".format(len(result))
+            
+            # Create data for bar chart (list of tuples)
+            chart_data = []
+            
+            for i, item in enumerate(result, 1):
+                score_pct = item['score'] * 100
+                output_text += f"{i}. **{item['label']}** (ID: {item['id']}): {score_pct:.2f}%\n"
+                chart_data.append((item['label'], item['score']))
+            
+            return output_text, chart_data
+        else:
+            return f"Error: {result.get('error', 'Unknown error')}", None
+            
+    except Exception as e:
+        return f"Error: {str(e)}", None
+
+def upsert_labels_admin(admin_token, new_items_json):
+    """
+    Admin function to add new labels.
+    """
+    if handler is None:
+        return "Error: Handler not initialized"
+        
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        # Parse the JSON input
+        items = json.loads(new_items_json) if new_items_json else []
+        
+        result = handler({
+            "inputs": {
+                "op": "upsert_labels",
+                "token": admin_token,
+                "items": items
+            }
+        })
+        
+        if result.get("status") == "ok":
+            return f"✅ Success! Added {result.get('added', 0)} new labels. Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        else:
+            return f"❌ Error: {result.get('detail', result.get('error', 'Unknown error'))}"
+            
+    except json.JSONDecodeError:
+        return "❌ Error: Invalid JSON format"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def reload_labels_admin(admin_token, version):
+    """
+    Admin function to reload a specific label version.
+    """
+    if handler is None:
+        return "Error: Handler not initialized"
+        
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        result = handler({
+            "inputs": {
+                "op": "reload_labels",
+                "token": admin_token,
+                "version": int(version) if version else 1
+            }
+        })
+        
+        if result.get("status") == "ok":
+            return f"✅ Labels reloaded successfully! Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("status") == "nochange":
+            return f"ℹ️ No change needed. Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        elif result.get("error") == "invalid_version":
+            return "❌ Error: Invalid version number"
+        else:
+            return f"❌ Error: {result.get('error', 'Unknown error')}"
+            
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def get_current_stats():
+    """
+    Get current label statistics.
+    """
+    if handler is None:
+        return "Handler not initialized"
+        
+    try:
+        num_labels = len(handler.class_ids) if hasattr(handler, 'class_ids') else 0
+        version = getattr(handler, 'labels_version', 1)
+        device = handler.device if hasattr(handler, 'device') else "unknown"
+        
+        stats = f"""
+**Current Statistics:**
+- Number of labels: {num_labels}
+- Labels version: {version}
+- Device: {device}
+- Model: MobileCLIP-B
+        """
+        
+        if hasattr(handler, 'class_names') and len(handler.class_names) > 0:
+            stats += f"\n- Sample labels: {', '.join(handler.class_names[:5])}"
+            if len(handler.class_names) > 5:
+                stats += "..."
+        
+        return stats
+    except Exception as e:
+        return f"Error getting stats: {str(e)}"
+
+# Create Gradio interface
+print("Creating Gradio interface...")
+with gr.Blocks(title="MobileCLIP Image Classifier") as demo:
+    gr.Markdown("""
+    # 🖼️ MobileCLIP-B Zero-Shot Image Classifier
+    
+    Upload an image to classify it using MobileCLIP-B model with dynamic label management.
+    """)
+    
+    with gr.Tab("🔍 Image Classification"):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    type="pil",
+                    label="Upload Image"
+                )
+                top_k_slider = gr.Slider(
+                    minimum=1,
+                    maximum=50,
+                    value=10,
+                    step=1,
+                    label="Number of top results to show"
+                )
+                classify_btn = gr.Button("🚀 Classify Image", variant="primary")
+            
+            with gr.Column():
+                output_text = gr.Markdown(label="Classification Results")
+                # Simplified bar chart using Dataframe
+                output_chart = gr.Dataframe(
+                    headers=["Label", "Confidence"],
+                    label="Classification Scores",
+                    interactive=False
+                )
+        
+        # Event handler for classification
+        classify_btn.click(
+            fn=classify_image,
+            inputs=[input_image, top_k_slider],
+            outputs=[output_text, output_chart]
+        )
+        
+        # Also trigger on image upload
+        input_image.change(
+            fn=classify_image,
+            inputs=[input_image, top_k_slider],
+            outputs=[output_text, output_chart]
+        )
+    
+    with gr.Tab("🔧 Admin Panel"):
+        gr.Markdown("""
+        ### Admin Functions
+        **Note:** Requires admin token (set via environment variable `ADMIN_TOKEN`)
+        """)
+        
+        with gr.Row():
+            admin_token_input = gr.Textbox(
+                label="Admin Token",
+                type="password",
+                placeholder="Enter admin token"
+            )
+        
+        with gr.Accordion("📊 Current Statistics", open=True):
+            stats_display = gr.Markdown(value=get_current_stats())
+            refresh_stats_btn = gr.Button("🔄 Refresh Stats")
+            refresh_stats_btn.click(
+                fn=get_current_stats,
+                inputs=[],
+                outputs=stats_display
+            )
+        
+        with gr.Accordion("➕ Add New Labels", open=False):
+            gr.Markdown("""
+            Add new labels by providing JSON array:
+            ```json
+            [
+                {"id": 100, "name": "new_object", "prompt": "a photo of a new_object"},
+                {"id": 101, "name": "another_object", "prompt": "a photo of another_object"}
+            ]
+            ```
+            """)
+            new_items_input = gr.Code(
+                label="New Items JSON",
+                language="json",
+                lines=5,
+                value='[\n  {"id": 100, "name": "example", "prompt": "a photo of example"}\n]'
+            )
+            upsert_btn = gr.Button("➕ Add Labels", variant="primary")
+            upsert_output = gr.Markdown()
+            
+            upsert_btn.click(
+                fn=upsert_labels_admin,
+                inputs=[admin_token_input, new_items_input],
+                outputs=upsert_output
+            )
+        
+        with gr.Accordion("🔄 Reload Label Version", open=False):
+            gr.Markdown("Reload labels from a specific version stored in the Hub")
+            version_input = gr.Number(
+                label="Version Number",
+                value=1,
+                precision=0
+            )
+            reload_btn = gr.Button("🔄 Reload Version", variant="primary")
+            reload_output = gr.Markdown()
+            
+            reload_btn.click(
+                fn=reload_labels_admin,
+                inputs=[admin_token_input, version_input],
+                outputs=reload_output
+            )
+    
+    with gr.Tab("ℹ️ About"):
+        gr.Markdown("""
+        ## About MobileCLIP-B Classifier
+        
+        This Space provides a web interface for Apple's MobileCLIP-B model, optimized for fast zero-shot image classification.
+        
+        ### Features:
+        - 🚀 **Fast inference**: < 30ms on GPU
+        - 🏷️ **Dynamic labels**: Add/update labels without redeployment
+        - 🔄 **Version control**: Track and reload label versions
+        - 📊 **Visual results**: Classification scores and confidence
+        
+        ### Environment Variables (set in Space Settings):
+        - `ADMIN_TOKEN`: Secret token for admin operations
+        - `HF_LABEL_REPO`: Hub repository for label storage
+        - `HF_WRITE_TOKEN`: Token with write permissions to label repo
+        - `HF_READ_TOKEN`: Token with read permissions (optional)
+        
+        ### Model Details:
+        - **Architecture**: MobileCLIP-B with MobileOne blocks
+        - **Text Encoder**: Transformer-based, 77 token context
+        - **Image Size**: 224x224
+        - **Embedding Dim**: 512
+        
+        ### License:
+        Model weights are licensed under Apple Sample Code License (ASCL).
+        """)
+
+print("Gradio interface created successfully!")
+
+if __name__ == "__main__":
+    print("Launching Gradio app...")
+    demo.launch()

handler.py

@@ -93,6 +93,24 @@ class EndpointHandler:
             ok = self._load_snapshot_from_hub_version(ver)
             return {"status": "ok" if ok else "nochange", "labels_version": getattr(self, "labels_version", 0)}
 
+        # Admin op: remove_labels
+        if op == "remove_labels":
+            if payload.get("token") != ADMIN_TOKEN:
+                return {"error": "unauthorized"}
+            ids_to_remove = set(payload.get("ids", []))
+            if not ids_to_remove:
+                return {"error": "no_ids_provided"}
+            
+            removed = self._remove_items(ids_to_remove)
+            if removed > 0:
+                new_ver = int(getattr(self, "labels_version", 1)) + 1
+                try:
+                    self._persist_snapshot_to_hub(new_ver)
+                    self.labels_version = new_ver
+                except Exception as e:
+                    return {"status": "error", "removed": removed, "detail": str(e)}
+            return {"status": "ok", "removed": removed, "labels_version": getattr(self, "labels_version", 1)}
+
         # Freshness guard (optional)
         min_ver = payload.get("min_labels_version")
         if isinstance(min_ver, int) and min_ver > getattr(self, "labels_version", 0):
@@ -151,6 +169,37 @@ class EndpointHandler:
             self._to_device()
             return len(batch)
 
+    def _remove_items(self, ids_to_remove):
+        if not ids_to_remove or not hasattr(self, "class_ids"):
+            return 0
+        with self._lock:
+            ids_to_remove = set(int(id) for id in ids_to_remove)
+            # Find indices to keep
+            indices_to_keep = []
+            removed_count = 0
+            for i, class_id in enumerate(self.class_ids):
+                if class_id not in ids_to_remove:
+                    indices_to_keep.append(i)
+                else:
+                    removed_count += 1
+            
+            if removed_count == 0:
+                return 0
+            
+            # Filter the tensors and lists
+            if indices_to_keep:
+                self.text_features_cpu = self.text_features_cpu[indices_to_keep].contiguous()
+                self.class_ids = [self.class_ids[i] for i in indices_to_keep]
+                self.class_names = [self.class_names[i] for i in indices_to_keep]
+            else:
+                # All items removed, reset to empty
+                self.text_features_cpu = torch.empty(0, self.text_features_cpu.shape[1])
+                self.class_ids = []
+                self.class_names = []
+            
+            self._to_device()
+            return removed_count
+
     def _persist_snapshot_to_hub(self, version: int):
         if not HF_LABEL_REPO:
             raise RuntimeError("HF_LABEL_REPO not set")