Update app.py

app.py

@@ -4,16 +4,16 @@ from transformers import AutoTokenizer
 import numpy as np
 
 # Initialize models
-context_model_file = "./bart-large-mnli.onnx"
+context_model_file = "./bart-base-mnli.onnx"  # Using bart-base-mnli for faster inference
 translation_model_file = "./model.onnx"
 
 # Create inference sessions for both models
 context_session = ort.InferenceSession(context_model_file)
 translation_session = ort.InferenceSession(translation_model_file)
 
-# Load tokenizer for translation model and for context model
-tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")
-context_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-mnli")
+# Load tokenizers for context and translation models
+context_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base-mnli")
+translation_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")
 
 labels = [
     "general", "pharma", "legal", "technical", "UI", "user interface", "medicine", 
@@ -23,51 +23,71 @@ labels = [
     "finance", "sports", "education", "politics", "economics", "art", "history",
     "music", "gaming", "aerospace", "engineering", "robotics", "travel", "tourism",
     "healthcare", "psychology", "environment", "fashion", "design", "real estate",
-    "retail", "news", "entertainment", "social media","automotive", "machine learning",
-    "cryptocurrency","blockchain","philosophy","anthropology","archaeology","data science"
+    "retail", "news", "entertainment", "social media", "automotive", "machine learning",
+    "cryptocurrency", "blockchain", "philosophy", "anthropology", "archaeology", "data science"
 ]
 
-def detect_context(input_text, top_n=3, score_threshold=0.05):
+def softmax_with_temperature(logits, temperature=1.0):
+    exp_logits = np.exp(logits / temperature)
+    return exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
+
+def detect_context(input_text, temperature=2.0, score_threshold=0.05):
     # Tokenize input text
     inputs = context_tokenizer(input_text, return_tensors="np", padding=True, truncation=True, max_length=512)
-    
     input_ids = inputs["input_ids"].astype(np.int64)
     attention_mask = inputs["attention_mask"].astype(np.int64)
 
+    # Debugging: Check tokenized input
+    print(f"Tokenized Input IDs: {input_ids}")
+    print(f"Tokenized Attention Mask: {attention_mask}")
+
     # Run inference with the ONNX context model
     outputs = context_session.run(None, {
         "input_ids": input_ids,
         "attention_mask": attention_mask
     })
 
-    scores = outputs[0][0]  # Assuming batch size 1; take the first set of logits
+    # Debugging: Check output shape
+    print(f"Logits shape: {outputs[0].shape}")  # Expected: (batch_size, num_labels)
+
+    logits = outputs[0][0]  # Assuming batch size 1; take the first set of logits
+
+    # Debugging: Print raw logits
+    print(f"Raw logits: {logits}")
+
+    # Apply softmax with temperature
+    scores = softmax_with_temperature(logits, temperature=temperature)
+
+    # Debugging: Print scores
+    print(f"Scores with softmax: {scores}")
 
     # Pair labels with scores
     label_scores = [(label, score) for label, score in zip(labels, scores)]
 
+    # Debugging: Print all label scores
+    print(f"All label scores: {label_scores}")
+
     # Sort by scores in descending order
     sorted_labels = sorted(label_scores, key=lambda x: x[1], reverse=True)
 
-    # Filter by threshold and return top_n contexts
+    # Filter by threshold
     filtered_labels = [label for label, score in sorted_labels if score > score_threshold]
-    top_contexts = filtered_labels[:top_n]
-
-    print(f"All scores: {label_scores}")  # Debugging: Print all scores
-    print(f"Selected contexts: {top_contexts}")  # Debugging: Print selected contexts
 
-    return top_contexts if top_contexts else ["general"]
+    # Debugging: Print filtered labels
+    print(f"Filtered labels: {filtered_labels}")
 
+    # Default to "general" if no valid context is found
+    return filtered_labels if filtered_labels else ["general"]
 
 def translate_text(input_text):
-    tokenized_input = tokenizer(
+    tokenized_input = translation_tokenizer(
         input_text, return_tensors="np",
         padding=True, truncation=True, max_length=512
     )
-    
     input_ids = tokenized_input["input_ids"].astype(np.int64)
     attention_mask = tokenized_input["attention_mask"].astype(np.int64)
 
-    decoder_start_token_id = tokenizer.cls_token_id or tokenizer.pad_token_id
+    decoder_start_token_id = translation_tokenizer.cls_token_id or translation_tokenizer.pad_token_id
     decoder_input_ids = np.array([[decoder_start_token_id]], dtype=np.int64)
 
     for _ in range(512):
@@ -86,10 +106,10 @@ def translate_text(input_text):
             [decoder_input_ids, np.array([[next_token_id]], dtype=np.int64)], axis=1
         )
 
-        if next_token_id == tokenizer.eos_token_id:
+        if next_token_id == translation_tokenizer.eos_token_id:
             break
 
-    return tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
+    return translation_tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
 
 def process_request(input_text):
     context = detect_context(input_text)
@@ -101,4 +121,4 @@ gr.Interface(
     inputs="text",
     outputs="text",
     live=True
-).launch()
+).launch()