Update app.py

app.py

@@ -27,31 +27,36 @@ labels = [
     "cryptocurrency","blockchain","philosophy","anthropology","archaeology","data science"
 ]
 
-def detect_context(input_text):
+def detect_context(input_text, top_n=3, score_threshold=0.05):
     # Tokenize input text
     inputs = context_tokenizer(input_text, return_tensors="np", padding=True, truncation=True, max_length=512)
     
-    # Prepare input for ONNX model
     input_ids = inputs["input_ids"].astype(np.int64)
     attention_mask = inputs["attention_mask"].astype(np.int64)
-    
-    # Run inference with context model
+
+    # Run inference with the ONNX context model
     outputs = context_session.run(None, {
         "input_ids": input_ids,
         "attention_mask": attention_mask
-        })
-    
-    # Assuming the output is logits for each label
-    scores = outputs[0]
-    
-    # Get the top label based on scores
-    top_label_index = np.argmax(scores, axis=1)[0]
-    
-    # Map index to label
-    detected_context = labels[top_label_index]
-    
-    print(detected_context)
-    return detected_context
+    })
+
+    scores = outputs[0][0]  # Assuming batch size 1; take the first set of logits
+
+    # Pair labels with scores
+    label_scores = [(label, score) for label, score in zip(labels, 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
+    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"]
+
 
 def translate_text(input_text):
     tokenized_input = tokenizer(