Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModel
 from sklearn.metrics.pairwise import cosine_similarity
 import torch
+import numpy as np
 from gradio_client import Client
 from functools import lru_cache
 
@@ -41,8 +42,13 @@ def precompute_label_embeddings():
 
 label_embeddings = precompute_label_embeddings()
 
+# Softmax function to convert scores to probabilities
+def softmax(x):
+    exp_x = np.exp(x - np.max(x))  # Subtract max for numerical stability
+    return exp_x / exp_x.sum()
+
 # Function to detect context
-def detect_context(input_text, fallback_threshold=0.5):  # Lowered threshold for debugging
+def detect_context(input_text, top_n=3):
     # Encode the input text
     inputs = tokenizer([input_text], padding=True, truncation=True, return_tensors="pt")
     with torch.no_grad():
@@ -52,19 +58,19 @@ def detect_context(input_text, fallback_threshold=0.5):  # Lowered threshold for
     # Compute cosine similarities
     similarities = cosine_similarity(input_embedding, label_embeddings)[0]
 
-    # Debugging: Print all labels and their similarity scores
-    print("Debug: Similarity scores for all labels:")
-    for label, score in zip(labels, similarities):
-        print(f"{label}: {score:.4f}")
+    # Apply softmax to convert similarities to probabilities
+    probabilities = softmax(similarities)
+
+    # Pair each label with its probability
+    label_probabilities = list(zip(labels, probabilities))
 
-    # Filter contexts with confidence >= fallback_threshold
-    high_confidence_contexts = [(labels[i], score) for i, score in enumerate(similarities) if score >= fallback_threshold]
+    # Sort by probability in descending order
+    label_probabilities.sort(key=lambda x: x[1], reverse=True)
 
-    # If no contexts meet the threshold, include "general" as a fallback
-    if not high_confidence_contexts:
-        high_confidence_contexts = [("general", 1.0)]  # Assign a default score of 1.0 for "general"
+    # Select the top N contexts
+    top_contexts = label_probabilities[:top_n]
 
-    return high_confidence_contexts
+    return top_contexts
 
 # Translation client
 translation_client = Client("Frenchizer/space_7")
@@ -81,7 +87,7 @@ def process_request(input_text):
     context_results = detect_context(input_text)
 
     # Step 3: Print the list of high-confidence contexts
-    print("High-confidence contexts:", context_results)
+    print("Detected Contexts (Top 3):", context_results)
 
     # Return the translation and contexts
     return translation, context_results
@@ -90,9 +96,9 @@ def process_request(input_text):
 def gradio_interface(input_text):
     translation, contexts = process_request(input_text)
     # Format the output
-    output = f"Translation: {translation}\n\nDetected Contexts:\n"
+    output = f"Translation: {translation}\n\nDetected Contexts (Top 3):\n"
     for context, score in contexts:
-        output += f"- {context} (confidence: {score:.2f})\n"
+        output += f"- {context} (confidence: {score:.4f})\n"
     return output.strip()
 
 # Create the Gradio interface