Update app.py

app.py

@@ -4,47 +4,64 @@ import torch
 
 # Load model and tokenizer
 model_name = "cross-encoder/ms-marco-MiniLM-L-12-v2"
-
-print("Loading model and tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 model.eval()  # Set model to evaluation mode
-print("Model and tokenizer loaded successfully.")
 
-# Function to compute relevance score
-def get_relevance_score(query, paragraph):
+# Function to get relevance score and relevant excerpt based on attention scores
+def get_relevance_score_and_excerpt(query, paragraph):
     if not query.strip() or not paragraph.strip():
-        return "Please provide both a query and a document paragraph."
+        return "Please provide both a query and a document paragraph.", ""
+    
+    # Tokenize the input
+    inputs = tokenizer(query, paragraph, return_tensors="pt", truncation=True, padding=True, return_attention_mask=True)
     
-    print(f"Received inputs -> Query: {query}, Paragraph: {paragraph}")
+    with torch.no_grad():
+        output = model(**inputs, output_attentions=True)  # Get attention scores
 
-    # Tokenize inputs
-    inputs = tokenizer(query, paragraph, return_tensors="pt", truncation=True, padding=True)
+    # Extract logits and calculate relevance score
+    logit = output.logits.squeeze().item()
+    relevance_score = torch.sigmoid(torch.tensor(logit)).item()
 
-    # Perform inference without gradient tracking
-    with torch.no_grad():
-        score = model(**inputs).logits.squeeze().item()
+    # Extract attention scores (use the last attention layer)
+    attention = output.attentions[-1]  # Shape: (batch_size, num_heads, seq_len, seq_len)
+
+    # Average across attention heads to get token importance
+    attention_scores = attention.mean(dim=1).squeeze(0)  # Shape: (seq_len, seq_len)
     
-    print(f"Calculated score: {score}")
-    return round(score, 4)
+    # Focus on the paragraph part only (ignore query tokens)
+    input_tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze())
+    query_length = len(tokenizer.tokenize(query))
+    
+    # Extract attention for the paragraph tokens only
+    paragraph_tokens = input_tokens[query_length + 2 : -1]  # Skip query and special tokens like [SEP]
+    paragraph_attention = attention_scores[query_length + 2 : -1, query_length + 2 : -1].mean(dim=0)
+
+    # Get the top tokens with highest attention scores
+    top_token_indices = torch.argsort(paragraph_attention, descending=True)[:5]  # Top 5 tokens
+    highlighted_tokens = [paragraph_tokens[i] for i in top_token_indices]
 
+    # Reconstruct the excerpt from top attention tokens
+    excerpt = tokenizer.convert_tokens_to_string(highlighted_tokens)
 
-def test_function(query, paragraph):
-    return f"Received query: {query}, paragraph: {paragraph}"
+    return round(relevance_score, 4), excerpt
 
 # Define Gradio interface
 interface = gr.Interface(
-    fn=get_relevance_score,
+    fn=get_relevance_score_and_excerpt,
     inputs=[
         gr.Textbox(label="Query", placeholder="Enter your search query..."),
         gr.Textbox(label="Document Paragraph", placeholder="Enter a paragraph to match...")
     ],
-    outputs=gr.Textbox(label="Relevance Score"),
-    title="Cross-Encoder Relevance Scoring",
-    description="Enter a query and a document paragraph to get a relevance score using the MS MARCO MiniLM L-12 v2 model."
+    outputs=[
+        gr.Textbox(label="Relevance Score"),
+        gr.Textbox(label="Most Relevant Excerpt")
+    ],
+    title="Cross-Encoder Relevance Scoring with Attention-Based Excerpt Extraction",
+    description="Enter a query and a document paragraph to get a relevance score and a relevant excerpt using attention scores.",
+    allow_flagging="never",
+    live=True
 )
 
-
 if __name__ == "__main__":
-    print("Launching Gradio app...")
     interface.launch()