Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
 # Define model names
 models = {
+    "gte base (gender v3.1)": "breadlicker45/gte-gender-v3.1-test",
     "ModernBERT Large (gender v3)": "breadlicker45/modernbert-gender-v3-test",
     "ModernBERT Large (gender v2)": "breadlicker45/modernbert-gender-v2",
     "ModernBERT Base (gender)": "breadlicker45/ModernBERT-base-gender",
@@ -10,8 +12,6 @@ models = {
 }
 
 # Define the mapping for user-friendly labels
-# Note: Transformers pipelines often output 'LABEL_0', 'LABEL_1'.
-# We handle potential variations like just '0', '1'.
 label_map = {
     "LABEL_0": "Male (0)",
     "0": "Male (0)",
@@ -19,27 +19,74 @@ label_map = {
     "1": "Female (1)"
 }
 
-# Function to load the selected model and classify text
+# A cache to store loaded models/pipelines to speed up subsequent requests
+model_cache = {}
+
+# Determine the device to run on (GPU if available, otherwise CPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+
+# The main classification function, now handles both model types
 def classify_text(model_name, text):
     try:
-        classifier = pipeline("text-classification", model=models[model_name], top_k=None)
-        predictions = classifier(text)
-
-        # Process predictions to use friendly labels
         processed_results = {}
-        if predictions and isinstance(predictions, list) and predictions[0]:
-             # predictions[0] should be a list of label dicts like [{'label': 'LABEL_1', 'score': 0.9...}, ...]
-            for pred in predictions[0]:
-                raw_label = pred["label"]
-                score = pred["score"]
-                # Use the map to get a friendly name, fallback to the raw label if not found
-                friendly_label = label_map.get(raw_label, raw_label)
-                processed_results[friendly_label] = score
+        model_id = models[model_name]
+
+        # --- SPECIAL HANDLING FOR THE GTE MODEL ---
+        if "gte-gender" in model_id:
+            # Check if model/tokenizer is already in our cache
+            if model_id not in model_cache:
+                print(f"Loading GTE model and tokenizer manually: {model_id}...")
+                tokenizer = AutoTokenizer.from_pretrained(model_id)
+                model = AutoModelForSequenceClassification.from_pretrained(model_id, trust_remote_code=True).to(device)
+                model_cache[model_id] = (model, tokenizer) # Cache both
+            
+            model, tokenizer = model_cache[model_id]
+
+            # Tokenize the input text and move to the correct device
+            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+
+            # Get model predictions
+            with torch.no_grad():
+                logits = model(**inputs).logits
+            
+            # Convert logits to probabilities using softmax
+            probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
+            
+            # Format results to match the pipeline's output style
+            processed_results[label_map["LABEL_0"]] = probabilities[0].item()
+            processed_results[label_map["LABEL_1"]] = probabilities[1].item()
+
+        # --- STANDARD HANDLING FOR PIPELINE-COMPATIBLE MODELS ---
+        else:
+            # Check if the pipeline is already in our cache
+            if model_id not in model_cache:
+                print(f"Loading pipeline for model: {model_id}...")
+                # Load and cache the pipeline
+                model_cache[model_id] = pipeline(
+                    "text-classification", 
+                    model=model_id, 
+                    top_k=None,
+                    device=device # Use the determined device
+                )
+            
+            classifier = model_cache[model_id]
+            predictions = classifier(text)
+
+            # Process predictions to use friendly labels
+            if predictions and isinstance(predictions, list) and predictions[0]:
+                for pred in predictions[0]:
+                    raw_label = pred["label"]
+                    score = pred["score"]
+                    friendly_label = label_map.get(raw_label, raw_label)
+                    processed_results[friendly_label] = score
+        
         return processed_results
+
     except Exception as e:
-        # Handle potential errors during model loading or inference
         print(f"Error: {e}")
-        # Return an error message suitable for gr.Label
+        # Return an error message suitable for gr.Label or gr.JSON
         return {"Error": f"Failed to process: {e}"}
 
 
@@ -50,20 +97,22 @@ interface = gr.Interface(
         gr.Dropdown(
             list(models.keys()),
             label="Select Model",
-            value="ModernBERT Large (gender)" # Default model
+            value="gte base (gender v3.1)" # Default model
         ),
         gr.Textbox(
             lines=2,
-            placeholder="Enter text to classify for perceived gender...", # Corrected placeholder
-            value="This is an example sentence." # Changed example text
+            placeholder="Enter text to classify for perceived gender...",
+            value="This is an example sentence."
         )
     ],
-    # The gr.Label component works well for showing classification scores
-    outputs=gr.Label(num_top_classes=2), # Show both classes explicitly
-    title="ModernBERT Gender Classifier",
-    description="Select a model and enter a sentence to see the perceived gender classification (Male=0, Female=1) and confidence scores. Note: Text-based gender classification can be unreliable and reflect societal biases.", # Updated description
+    # Since we now consistently return a dictionary of {label: score},
+    # we can go back to using the nicer-looking gr.Label component!
+    outputs=gr.Label(num_top_classes=2, label="Classification Results"),
+    title="ModernBERT & GTE Gender Classifier",
+    description="Select a model and enter a sentence to see the perceived gender classification (Male=0, Female=1) and confidence scores. Note: Text-based gender classification can be unreliable and reflect societal biases.",
+    allow_flagging="never",
 )
 
 # Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()