Update app.py

app.py

@@ -8,6 +8,10 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 model.eval()
 
+# Sigmoid-based threshold adjustment function
+def calculate_threshold(base_relevance, min_threshold=0.02, max_threshold=0.5, k=10):
+    return min_threshold + (max_threshold - min_threshold) * (1 / (1 + torch.exp(-k * (base_relevance - 0.5))))
+
 # Function to compute relevance score and dynamically adjust threshold
 def get_relevance_score_and_excerpt(query, paragraph, threshold_weight):
     if not query.strip() or not paragraph.strip():
@@ -23,8 +27,8 @@ def get_relevance_score_and_excerpt(query, paragraph, threshold_weight):
     logit = output.logits.squeeze().item()
     base_relevance_score = torch.sigmoid(torch.tensor(logit)).item()
 
-    # Dynamically adjust the attention threshold based on user weight (no relevance score influence)
-    dynamic_threshold = max(0.02, threshold_weight)
+    # Compute dynamic threshold using sigmoid-based adjustment
+    dynamic_threshold = calculate_threshold(base_relevance_score) * threshold_weight
 
     # Extract attention scores (last layer)
     attention = output.attentions[-1]
@@ -66,14 +70,14 @@ interface = gr.Interface(
     inputs=[
         gr.Textbox(label="Query", placeholder="Enter your search query..."),
         gr.Textbox(label="Document Paragraph", placeholder="Enter a paragraph to match..."),
-        gr.Slider(minimum=0.02, maximum=0.5, value=0.1, step=0.01, label="Attention Threshold")
+        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Threshold Weight")
     ],
     outputs=[
         gr.Textbox(label="Relevance Score"),
         gr.HTML(label="Highlighted Document Paragraph")
     ],
     title="Cross-Encoder Attention Highlighting",
-    description="Adjust the attention threshold to control token highlighting sensitivity.",
+    description="Adjust the attention threshold weight to control token highlighting sensitivity.",
     allow_flagging="never",
     live=True
 )