Update app.py

app.py

@@ -23,14 +23,26 @@ labels = [
     "automotive", "blockchain", "biology", "chemistry", 
     "cryptocurrency", "data science", "design", "e-commerce",
     "education", "engineering", "entertainment", "environment",
-    "fashion", "finance", "food commerce", "general",
-    "gaming", "healthcare", "history", "html",
-    "information technology", "IT", "keywords", "legal",
-    "literature", "machine learning", "marketing", "medicine",
-    "music", "personal development", "philosophy", "physics",
-    "politics", "poetry", "programming", "real estate", "retail", 
-    "robotics", "slang", "social media", "speech", "sports",
-    "sustained", "technical", "theater", "tourism", "travel"
+    "fashion", "finance", "food commerce", "gaming",
+    "healthcare", "history", "information technology", 
+    "legal", "machine learning", "marketing", "medicine",
+    "music", "philosophy", "physics", "politics", "real estate", "retail", 
+    "robotics", "social media", "sports", "technical",
+    "tourism", "travel"
+]
+
+tones = [
+    "formal", "positive", "negative", "poetic", "polite", "subtle", "casual", "neutral",
+    "informal", "pompous", "sustained", "rude", "sustained",
+    "sophisticated", "playful", "serious", "friendly"
+]
+
+styles = [
+    "poetry", "novel", "theater", "slang", "speech", "keywords", "html", "programming"
+]
+
+gender_number = [
+    "masculine singular", "masculine plural", "feminine singular", "feminine plural"
 ]
 
 @lru_cache(maxsize=1)
@@ -48,7 +60,7 @@ def softmax(x):
     return exp_x / exp_x.sum()
 
 # Function to detect context
-def detect_context(input_text, top_n=3):
+def detect_context(input_text, threshold=0.022):
     # Encode the input text
     inputs = tokenizer([input_text], padding=True, truncation=True, return_tensors="pt")
     with torch.no_grad():
@@ -64,14 +76,14 @@ def detect_context(input_text, top_n=3):
     # Pair each label with its probability
     label_probabilities = list(zip(labels, probabilities))
 
-    # Sort by probability in descending order
-    label_probabilities.sort(key=lambda x: x[1], reverse=True)
+    # Filter contexts with confidence >= threshold
+    high_confidence_contexts = [(label, score) for label, score in label_probabilities if score >= threshold]
 
-    # Select the top N contexts
-    top_contexts = label_probabilities[:top_n]
+    # If no contexts meet the threshold, default to "general"
+    if not high_confidence_contexts:
+        high_confidence_contexts = [("general", 1.0)]  # Assign a default score of 1.0 for "general"
 
-    # Return both the top N contexts and all context scores
-    return top_contexts, label_probabilities
+    return high_confidence_contexts
 
 # Translation client
 translation_client = Client("Frenchizer/space_7")
@@ -85,27 +97,21 @@ def process_request(input_text):
     translation = translate_text(input_text)
 
     # Step 2: Detect context
-    top_contexts, all_contexts = detect_context(input_text)
+    context_results = detect_context(input_text)
 
-    # Step 3: Print the list of high-confidence contexts and all context scores
-    print("Detected Contexts (Top 3):", top_contexts)
-    print("All Context Scores:")
-    for context, score in all_contexts:
-        print(f"- {context}: {score:.4f}")
+    # Step 3: Print the list of high-confidence contexts
+    print("High-confidence contexts (score >= 0.022):", context_results)
 
     # Return the translation and contexts
-    return translation, top_contexts, all_contexts
+    return translation, context_results
 
 # Gradio interface
 def gradio_interface(input_text):
-    translation, top_contexts, all_contexts = process_request(input_text)
+    translation, contexts = process_request(input_text)
     # Format the output
-    output = f"Translation: {translation}\n\nDetected Contexts (Top 3):\n"
-    for context, score in top_contexts:
+    output = f"Translation: {translation}\n\nDetected Contexts (score >= 0.022):\n"
+    for context, score in contexts:
         output += f"- {context} (confidence: {score:.4f})\n"
-    output += "\nAll Context Scores:\n"
-    for context, score in all_contexts:
-        output += f"- {context}: {score:.4f}\n"
     return output.strip()
 
 # Create the Gradio interface