Update app.py

app.py

@@ -14,7 +14,7 @@ def get_relevance_score_and_excerpt(query, 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)
+    inputs = tokenizer(query, paragraph, return_tensors="pt", truncation=True, padding=True)
     
     with torch.no_grad():
         output = model(**inputs, output_attentions=True)  # Get attention scores
@@ -23,31 +23,37 @@ def get_relevance_score_and_excerpt(query, paragraph):
     logit = output.logits.squeeze().item()
     relevance_score = torch.sigmoid(torch.tensor(logit)).item()
 
-    # Extract attention scores (use the last attention layer)
+    # Extract attention scores (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)
-    
-    # Tokenize query and paragraph separately to get lengths
+    # Average attention across heads and batch dimension
+    attention_scores = attention.mean(dim=1).mean(dim=0)  # Shape: (seq_len, seq_len)
+
+    # Get tokenized query and paragraph separately
     query_tokens = tokenizer.tokenize(query)
     paragraph_tokens = tokenizer.tokenize(paragraph)
 
-    # Extract only paragraph-related attention scores
-    query_length = len(query_tokens) + 2  # +2 for special tokens like [CLS] and [SEP]
-    para_start_idx = query_length
+    query_len = len(query_tokens) + 2  # +2 for special tokens [CLS] and [SEP]
+    para_start_idx = query_len
     para_end_idx = len(inputs["input_ids"][0]) - 1  # Ignore final [SEP] token
 
-    para_attention_scores = attention_scores[0, para_start_idx:para_end_idx].mean(dim=0)
+    # Handle potential indexing issues
+    if para_end_idx <= para_start_idx:
+        return round(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(relevance_score, 4), "No relevant tokens extracted."
 
     # Get indices of top-k attended tokens while preserving order
-    top_k = min(5, len(paragraph_tokens))  # Extract top 5 tokens or fewer if short
-    top_indices = para_attention_scores.argsort(descending=True)[:top_k].sort().values  # Sort to preserve order
+    top_k = min(5, para_attention_scores.size(0))  # Ensure top-k does not exceed available tokens
+    top_indices = para_attention_scores.topk(top_k).indices.sort().values  # Sort indices to preserve order
 
-    # Extract top tokens based on original order
-    highlighted_tokens = [paragraph_tokens[i] for i in top_indices]
+    # Extract highlighted tokens from the paragraph
+    highlighted_tokens = [paragraph_tokens[i] for i in top_indices.tolist()]
 
-    # Reconstruct the excerpt from ordered tokens
+    # Convert tokens back to a readable string
     excerpt = tokenizer.convert_tokens_to_string(highlighted_tokens)
 
     return round(relevance_score, 4), excerpt