Update app.py

app.py

@@ -1,64 +1,44 @@
 import gradio as gr
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 # 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)
-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):
-    base_relevance_tensor = torch.tensor(base_relevance)  # Ensure input is a tensor
-    threshold = min_threshold + (max_threshold - min_threshold) * (
-        1 / (1 + torch.exp(-k * (base_relevance_tensor - 0.5)))
-    )
-    return threshold.item()  # Convert tensor back to float for use in other functions
-    
-# 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.", ""
+# Set model to evaluation mode
+model.eval()
 
-    # Tokenize the input
-    inputs = tokenizer(query, paragraph, return_tensors="pt", truncation=True, padding=True)
+# 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"]
     
+    # Get model outputs with attentions
     with torch.no_grad():
-        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()
-
-    # 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]
-    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), "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), "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
+        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])
     highlighted_text = ""
-    for idx, token in enumerate(paragraph_tokens):
+    for idx, token in enumerate(tokens):
         if idx in relevant_indices:
             highlighted_text += f"<b>{token}</b> "
         else:
@@ -66,25 +46,25 @@ def get_relevance_score_and_excerpt(query, paragraph, threshold_weight):
 
     highlighted_text = tokenizer.convert_tokens_to_string(highlighted_text.split())
 
-    return round(base_relevance_score, 4), highlighted_text
+    # Print values to debug
+    print(f"Relevance Score: {relevance_score}")
+    print(f"Dynamic Threshold: {dynamic_threshold}")
 
-# Define Gradio interface with a slider for threshold adjustment
-interface = gr.Interface(
-    fn=get_relevance_score_and_excerpt,
+    return relevance_score, dynamic_threshold, highlighted_text
+
+# Create Gradio interface with a slider for threshold adjustment weight
+iface = gr.Interface(
+    fn=process_text,
     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.5, maximum=2.0, value=1.0, step=0.1, label="Threshold Weight")
+        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"),
     ],
     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 weight to control token highlighting sensitivity.",
-    allow_flagging="never",
-    live=True
+    ]
 )
 
-if __name__ == "__main__":
-    interface.launch()
+iface.launch()