Update app.py

app.py

@@ -1,44 +1,56 @@
 import gradio as gr
-import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
 
 # Load model and tokenizer
 model_name = "cross-encoder/ms-marco-MiniLM-L-12-v2"
-model = AutoModelForSequenceClassification.from_pretrained(model_name, output_attentions=True)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Set model to evaluation mode
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 model.eval()
 
-# Function to compute relevance and highlight relevant tokens
-def process_text(query, document, weight):
-    # Tokenize input
-    inputs = tokenizer(query, document, return_tensors="pt", truncation=True, padding=True)
-    input_ids = inputs["input_ids"]
+# 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():
+        return "Please provide both a query and a document paragraph.", "", ""
+
+    # Tokenize the input
+    inputs = tokenizer(query, paragraph, return_tensors="pt", truncation=True, padding=True)
     
-    # Get model outputs with attentions
     with torch.no_grad():
-        outputs = model(**inputs, output_attentions=True)
-        relevance_score = torch.sigmoid(outputs.logits).item()  # Convert logits to relevance score
-        attentions = outputs.attentions[-1].squeeze(0).mean(0)  # Mean attention across heads
-
-    # Calculate dynamic threshold using sigmoid function
-    def calculate_threshold(base_relevance, min_threshold=0.0, max_threshold=0.5, k=10):
-        base_relevance_tensor = torch.tensor(base_relevance)
-        threshold = min_threshold + (max_threshold - min_threshold) * (
-            1 / (1 + torch.exp(-k * (base_relevance_tensor - 0.5)))
-        )
-        return threshold.item()
-
-    dynamic_threshold = calculate_threshold(relevance_score) * weight
-
-    # Extract important tokens based on attention scores
-    relevant_indices = (attentions > dynamic_threshold).nonzero(as_tuple=True)[0].tolist()
-    
-    # Highlight tokens in the original order, using HTML bold tags
-    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+        output = model(**inputs, output_attentions=True)
+
+    # Extract logits and calculate base relevance score
+    logit = output.logits.squeeze().item()
+    base_relevance_score = torch.sigmoid(torch.tensor(logit)).item()
+
+    # Dynamically adjust the attention threshold based on user weight
+    dynamic_threshold = max(0.02, threshold_weight)
+
+    # Extract attention scores (last layer)
+    attention = output.attentions[-1]
+    attention_scores = attention.mean(dim=1).mean(dim=0)
+
+    query_tokens = tokenizer.tokenize(query)
+    paragraph_tokens = tokenizer.tokenize(paragraph)
+
+    query_len = len(query_tokens) + 2  # +2 for special tokens [CLS] and first [SEP]
+    para_start_idx = query_len
+    para_end_idx = len(inputs["input_ids"][0]) - 1
+
+    if para_end_idx <= para_start_idx:
+        return round(base_relevance_score, 4), round(dynamic_threshold, 4), "No relevant tokens extracted."
+
+    para_attention_scores = attention_scores[para_start_idx:para_end_idx, para_start_idx:para_end_idx].mean(dim=0)
+
+    if para_attention_scores.numel() == 0:
+        return round(base_relevance_score, 4), round(dynamic_threshold, 4), "No relevant tokens extracted."
+
+    # Get indices of relevant tokens above dynamic threshold
+    relevant_indices = (para_attention_scores > dynamic_threshold).nonzero(as_tuple=True)[0].tolist()
+
+    # Reconstruct paragraph with bolded relevant tokens using HTML tags
     highlighted_text = ""
-    for idx, token in enumerate(tokens):
+    for idx, token in enumerate(paragraph_tokens):
         if idx in relevant_indices:
             highlighted_text += f"<b>{token}</b> "
         else:
@@ -46,25 +58,26 @@ def process_text(query, document, weight):
 
     highlighted_text = tokenizer.convert_tokens_to_string(highlighted_text.split())
 
-    # Print values to debug
-    print(f"Relevance Score: {relevance_score}")
-    print(f"Dynamic Threshold: {dynamic_threshold}")
+    return round(base_relevance_score, 4), round(dynamic_threshold, 4), highlighted_text
 
-    return relevance_score, dynamic_threshold, highlighted_text
-
-# Create Gradio interface with a slider for threshold adjustment weight
-iface = gr.Interface(
-    fn=process_text,
+# Define Gradio interface with a slider for threshold adjustment
+interface = gr.Interface(
+    fn=get_relevance_score_and_excerpt,
     inputs=[
-        gr.Textbox(label="Query"),
-        gr.Textbox(label="Document Paragraph"),
-        gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1.0, label="Threshold Weight"),
+        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")
     ],
     outputs=[
         gr.Textbox(label="Relevance Score"),
         gr.Textbox(label="Dynamic Threshold"),
         gr.HTML(label="Highlighted Document Paragraph")
-    ]
+    ],
+    title="Cross-Encoder Attention Highlighting",
+    description="Adjust the attention threshold to control token highlighting sensitivity.",
+    allow_flagging="never",
+    live=True
 )
 
-iface.launch()
+if __name__ == "__main__":
+    interface.launch()