added indexing for 1-2 documents at a time from cranfield and a viewing feature

app.py

@@ -5,6 +5,7 @@ import numpy as np
 from tqdm.auto import tqdm
 import os
 import ir_datasets
+import random # Added for random selection
 
 # --- Model Loading (Keep as is) ---
 tokenizer_splade = None
@@ -47,49 +48,77 @@ except Exception as e:
     print(f"Please ensure '{splade_doc_model_name}' is accessible (check Hugging Face Hub for potential agreements).")
 
 
-# --- Global Variables for Document Index ---
+# --- Global Variables for Document Index and Qrels ---
 document_representations = {} # Stores {doc_id: sparse_vector}
 document_texts = {}           # Stores {doc_id: doc_text}
+queries_texts = {}            # Stores {query_id: query_text}
+qrels_data = {}               # Stores {query_id: [{doc_id: str, relevance: int}, ...]}
 initial_doc_model_for_indexing = "SPLADE-cocondenser-distil" # Fixed for initial demo index
 
 
-# --- Load Cranfield Corpus using ir_datasets ---
-# Renamed function for clarity, but kept original name for call consistency
+# --- Load Cranfield Corpus, Queries, and Qrels using ir_datasets ---
 def load_cranfield_corpus_ir_datasets():
-    global document_texts
-    print("Loading Cranfield corpus using ir_datasets...")
+    global document_texts, queries_texts, qrels_data
+    print("Loading Cranfield corpus, queries, and qrels using ir_datasets...")
     try:
-        # --- IMPORTANT CHANGE: Loading 'cranfield' dataset ---
         dataset = ir_datasets.load("cranfield")
+
+        # Load documents
         for doc in tqdm(dataset.docs_iter(), desc="Loading Cranfield documents"):
             document_texts[doc.doc_id] = doc.text.strip()
         print(f"Loaded {len(document_texts)} documents from Cranfield corpus.")
+
+        # Load queries
+        for query in tqdm(dataset.queries_iter(), desc="Loading Cranfield queries"):
+            queries_texts[query.query_id] = query.text.strip()
+        print(f"Loaded {len(queries_texts)} queries from Cranfield corpus.")
+
+        # Load qrels
+        for qrel in tqdm(dataset.qrels_iter(), desc="Loading Cranfield qrels"):
+            if qrel.query_id not in qrels_data:
+                qrels_data[qrel.query_id] = []
+            qrels_data[qrel.query_id].append({"doc_id": qrel.doc_id, "relevance": qrel.relevance})
+        print(f"Loaded qrels for {len(qrels_data)} queries.")
+
     except Exception as e:
         print(f"Error loading Cranfield corpus with ir_datasets: {e}")
         print("Please ensure 'ir_datasets' is installed and your internet connection is stable.")
 
 
-# --- Helper function for lexical mask (Keep as is) ---
-def create_lexical_bow_mask(input_ids, vocab_size, tokenizer):
-    bow_mask = torch.zeros(vocab_size, device=input_ids.device)
-    meaningful_token_ids = []
-    for token_id in input_ids.squeeze().tolist():
-        if token_id not in [
-            tokenizer.pad_token_id,
-            tokenizer.cls_token_id,
-            tokenizer.sep_token_id,
-            tokenizer.mask_token_id,
-            tokenizer.unk_token_id
-        ]:
-            meaningful_token_ids.append(token_id)
+# --- Helper function for lexical mask (now handles batches) ---
+def create_lexical_bow_mask(input_ids_batch, vocab_size, tokenizer):
+    """
+    Creates a batch of lexical BOW masks.
+    input_ids_batch: torch.Tensor of shape (batch_size, sequence_length)
+    vocab_size: int, size of the tokenizer vocabulary
+    tokenizer: the tokenizer object
+    Returns: torch.Tensor of shape (batch_size, vocab_size)
+    """
+    batch_size = input_ids_batch.shape[0]
+    bow_masks = torch.zeros(batch_size, vocab_size, device=input_ids_batch.device)
+
+    for i in range(batch_size):
+        input_ids = input_ids_batch[i] # Get input_ids for the current item in the batch
+        meaningful_token_ids = []
+        for token_id in input_ids.tolist():
+            if token_id not in [
+                tokenizer.pad_token_id,
+                tokenizer.cls_token_id,
+                tokenizer.sep_token_id,
+                tokenizer.mask_token_id,
+                tokenizer.unk_token_id
+            ]:
+                meaningful_token_ids.append(token_id)
+                
+        if meaningful_token_ids:
+            # Apply mask to the current row in the batch
+            bow_masks[i, list(set(meaningful_token_ids))] = 1
             
-    if meaningful_token_ids:
-        bow_mask[list(set(meaningful_token_ids))] = 1
-            
-    return bow_mask.unsqueeze(0)
+    return bow_masks
 
 
 # --- Core Representation Functions (Return Formatted Strings - for Explorer Tab) ---
+# These functions still take single text input for the Explorer tab
 def get_splade_cocondenser_representation(text):
     if tokenizer_splade is None or model_splade is None:
         return "SPLADE-cocondenser-distil model is not loaded. Please check the console for loading errors."
@@ -104,7 +133,7 @@ def get_splade_cocondenser_representation(text):
         splade_vector = torch.max(
             torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1),
             dim=1
-        )[0].squeeze()
+        )[0].squeeze() # Squeeze is fine here as it's a single input
     else:
         return "Model output structure not as expected for SPLADE-cocondenser-distil. 'logits' not found."
 
@@ -151,15 +180,16 @@ def get_splade_lexical_representation(text):
         splade_vector = torch.max(
             torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1),
             dim=1
-        )[0].squeeze()
+        )[0].squeeze() # Squeeze is fine here
     else:
         return "Model output structure not as expected for SPLADE-v3-Lexical. 'logits' not found."
 
     # Always apply lexical mask for this model's specific behavior
     vocab_size = tokenizer_splade_lexical.vocab_size
+    # Call with unsqueezed input_ids for single sample processing
     bow_mask = create_lexical_bow_mask(
         inputs['input_ids'], vocab_size, tokenizer_splade_lexical
-    ).squeeze()
+    ).squeeze() # Squeeze back for single output
     splade_vector = splade_vector * bow_mask
 
     indices = torch.nonzero(splade_vector).squeeze().cpu().tolist()
@@ -202,12 +232,13 @@ def get_splade_doc_representation(text):
         output = model_splade_doc(**inputs)
 
     if not hasattr(output, "logits"):
-        return "SPLADE-v3-Doc model output structure not as expected. 'logits' not found."
+        return "Model output structure not as expected. 'logits' not found."
 
     vocab_size = tokenizer_splade_doc.vocab_size
+    # Call with unsqueezed input_ids for single sample processing
     binary_splade_vector = create_lexical_bow_mask(
         inputs['input_ids'], vocab_size, tokenizer_splade_doc
-    ).squeeze()
+    ).squeeze() # Squeeze back for single output
     
     indices = torch.nonzero(binary_splade_vector).squeeze().cpu().tolist()
     if not isinstance(indices, list):
@@ -253,44 +284,75 @@ def predict_representation_explorer(model_choice, text):
         return "Please select a model."
 
 
-# --- Internal Core Representation Functions (Return Raw Vectors - for Retrieval Tab) ---
-def get_splade_cocondenser_representation_internal(text, tokenizer, model):
+# --- Internal Core Representation Functions (now handle batches) ---
+def get_splade_cocondenser_representation_internal(texts, tokenizer, model):
+    """
+    Generates SPLADE representations for a batch of texts.
+    texts: list of strings
+    tokenizer: the tokenizer object
+    model: the SPLADE model
+    Returns: torch.Tensor of shape (batch_size, vocab_size) or None
+    """
     if tokenizer is None or model is None: return None
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    with torch.no_grad(): output = model(**inputs)
+
+    with torch.no_grad():
+        output = model(**inputs)
+
     if hasattr(output, 'logits'):
-        splade_vector = torch.max(torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1), dim=1)[0].squeeze()
-        return splade_vector
+        # torch.max(..., dim=1)[0] reduces along sequence_length dimension,
+        # resulting in (batch_size, vocab_size)
+        splade_vectors = torch.max(
+            torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1),
+            dim=1
+        )[0]
+        return splade_vectors
     else:
         print("Model output structure not as expected for SPLADE-cocondenser-distil. 'logits' not found.")
         return None
 
-def get_splade_lexical_representation_internal(text, tokenizer, model):
+def get_splade_lexical_representation_internal(texts, tokenizer, model):
+    """
+    Generates SPLADE-Lexical representations for a batch of texts.
+    texts: list of strings
+    tokenizer: the tokenizer object
+    model: the SPLADE-Lexical model
+    Returns: torch.Tensor of shape (batch_size, vocab_size) or None
+    """
     if tokenizer is None or model is None: return None
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
     with torch.no_grad(): output = model(**inputs)
     if hasattr(output, 'logits'):
-        splade_vector = torch.max(torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1), dim=1)[0].squeeze()
+        splade_vectors = torch.max(torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1), dim=1)[0]
         vocab_size = tokenizer.vocab_size
-        bow_mask = create_lexical_bow_mask(inputs['input_ids'], vocab_size, tokenizer).squeeze()
-        splade_vector = splade_vector * bow_mask
-        return splade_vector
+        # create_lexical_bow_mask now returns (batch_size, vocab_size)
+        bow_masks = create_lexical_bow_mask(inputs['input_ids'], vocab_size, tokenizer)
+        splade_vectors = splade_vectors * bow_masks # Element-wise multiplication, shapes (batch_size, vocab_size)
+        return splade_vectors
     else:
         print("Model output structure not as expected for SPLADE-v3-Lexical. 'logits' not found.")
         return None
 
-def get_splade_doc_representation_internal(text, tokenizer, model):
+def get_splade_doc_representation_internal(texts, tokenizer, model):
+    """
+    Generates SPLADE-Doc (binary) representations for a batch of texts.
+    texts: list of strings
+    tokenizer: the tokenizer object
+    model: the SPLADE-Doc model (not directly used for logits, but for device)
+    Returns: torch.Tensor of shape (batch_size, vocab_size) or None
+    """
     if tokenizer is None or model is None: return None
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()} # Ensure inputs are on the correct device
     vocab_size = tokenizer.vocab_size
-    binary_splade_vector = create_lexical_bow_mask(inputs['input_ids'], vocab_size, tokenizer).squeeze()
-    return binary_splade_vector
+    # create_lexical_bow_mask now returns (batch_size, vocab_size)
+    binary_splade_vectors = create_lexical_bow_mask(inputs['input_ids'], vocab_size, tokenizer)
+    return binary_splade_vectors
 
 
-# --- Document Indexing Function (for Retrieval Tab) ---
+# --- Document Indexing Function (now uses batching) ---
 def index_documents(doc_model_choice):
     global document_representations
     if document_representations:
@@ -328,14 +390,28 @@ def index_documents(doc_model_choice):
 
     print(f"Indexing documents using {doc_model_choice}...")
     
-    doc_items = list(document_texts.items())
+    doc_ids_list = list(document_texts.keys())
+    doc_texts_list = list(document_texts.values())
     
-    for doc_id, doc_text in tqdm(doc_items, desc="Indexing Documents"):
-        sparse_vector = representation_func_to_use(doc_text, tokenizer_to_use, model_to_use)
-        if sparse_vector is not None:
-            document_representations[doc_id] = sparse_vector.cpu()
+    # --- BATCH SIZE FOR INDEXING ---
+    batch_size = 32 # You can adjust this value based on memory and performance
+    
+    document_representations = {} # Ensure it's clear we're (re)building the index
+
+    # Iterate through documents in batches
+    for i in tqdm(range(0, len(doc_ids_list), batch_size), desc="Indexing Documents in Batches"):
+        batch_doc_ids = doc_ids_list[i:i + batch_size]
+        batch_doc_texts = doc_texts_list[i:i + batch_size]
+
+        sparse_vectors_batch = representation_func_to_use(batch_doc_texts, tokenizer_to_use, model_to_use)
+        
+        if sparse_vectors_batch is not None:
+            # sparse_vectors_batch will have shape (batch_size, vocab_size)
+            for j, doc_id in enumerate(batch_doc_ids):
+                # Store each document's vector
+                document_representations[doc_id] = sparse_vectors_batch[j].cpu()
         else:
-            print(f"Warning: Failed to get representation for doc_id {doc_id}")
+            print(f"Warning: Failed to get representation for a batch starting with doc_id {batch_doc_ids[0]}")
 
     print(f"Finished indexing {len(document_representations)} documents.")
     return True
@@ -349,25 +425,27 @@ def retrieve_documents(query_text, query_model_choice, indexed_doc_model_name, t
     query_tokenizer = None
     query_model = None
 
+    # These internal calls still use single text input for the query
     if query_model_choice == "SPLADE-cocondenser-distil (weighting and expansion)":
         query_tokenizer = tokenizer_splade
         query_model = model_splade
-        query_vector = get_splade_cocondenser_representation_internal(query_text, query_tokenizer, query_model)
+        query_vector = get_splade_cocondenser_representation_internal([query_text], query_tokenizer, query_model)
     elif query_model_choice == "SPLADE-v3-Lexical (weighting)":
         query_tokenizer = tokenizer_splade_lexical
         query_model = model_splade_lexical
-        query_vector = get_splade_lexical_representation_internal(query_text, query_tokenizer, query_model)
+        query_vector = get_splade_lexical_representation_internal([query_text], query_tokenizer, query_model)
     elif query_model_choice == "SPLADE-v3-Doc (binary)":
         query_tokenizer = tokenizer_splade_doc
         query_model = model_splade_doc
-        query_vector = get_splade_doc_representation_internal(query_text, query_tokenizer, query_model)
+        query_vector = get_splade_doc_representation_internal([query_text], query_tokenizer, query_model)
     else:
         return "Invalid query model choice.", []
 
     if query_vector is None:
         return "Failed to get query representation. Check console for model loading errors.", []
 
-    query_vector = query_vector.cpu()
+    # Since internal functions now return batches, take the first (and only) item for single query
+    query_vector = query_vector.squeeze(0).cpu()
 
     scores = {}
     for doc_id, doc_vec in document_representations.items():
@@ -396,9 +474,64 @@ def predict_retrieval_gradio(query_text, query_model_choice, selected_doc_model_
     formatted_output, _ = retrieve_documents(query_text, query_model_choice, initial_doc_model_for_indexing, top_k=5)
     return formatted_output
 
+# --- New function to get specific retrieval examples ---
+def get_specific_retrieval_examples():
+    if not queries_texts or not qrels_data or not document_texts:
+        return "Queries, qrels, or documents not loaded. Please check initial loading."
+
+    high_qrel_threshold = 3 # Relevance score of 3 or 4 for Cranfield is generally considered high
+    low_qrel_threshold = 1  # Relevance score of 0 or 1 for Cranfield is generally considered low
+
+    eligible_query_ids = []
+    for qid, qrels in qrels_data.items():
+        has_high_qrel = any(item['relevance'] >= high_qrel_threshold for item in qrels)
+        has_low_qrel = any(item['relevance'] <= low_qrel_threshold for item in qrels)
+        if has_high_qrel and has_low_qrel:
+            eligible_query_ids.append(qid)
+    
+    if not eligible_query_ids:
+        return "Could not find a query with both high and low relevance documents in the loaded qrels."
+
+    # Pick a random eligible query
+    random_query_id = random.choice(eligible_query_ids)
+    full_query_text = queries_texts.get(random_query_id, "Query text not found.")
+    query_snippet = full_query_text[:300] + "..." if len(full_query_text) > 300 else full_query_text
+    
+    qrels_for_query = qrels_data[random_query_id]
+
+    high_qrel_docs = [item for item in qrels_for_query if item['relevance'] >= high_qrel_threshold]
+    low_qrel_docs = [item for item in qrels_for_query if item['relevance'] <= low_qrel_threshold]
+
+    selected_high_doc_id = random.choice(high_qrel_docs)['doc_id'] if high_qrel_docs else None
+    selected_low_doc_id = random.choice(low_qrel_docs)['doc_id'] if low_qrel_docs else None
+
+    output_str = f"### Random Query Example\n\n"
+    output_str += f"**Query ID:** {random_query_id}\n"
+    output_str += f"**Query Snippet:** {query_snippet}\n\n" # Changed to snippet
+
+    if selected_high_doc_id:
+        full_doc_text = document_texts.get(selected_high_doc_id, "Document text not available.")
+        doc_snippet = full_doc_text[:500] + "..." if len(full_doc_text) > 500 else full_doc_text
+        output_str += f"### Highly Relevant Document (Qrel >= {high_qrel_threshold})\n"
+        output_str += f"**Document ID:** {selected_high_doc_id}\n"
+        output_str += f"**Document Snippet:** {doc_snippet}\n\n" # Changed to snippet
+    else:
+        output_str += "No highly relevant document found for this query.\n\n"
+
+    if selected_low_doc_id:
+        full_doc_text = document_texts.get(selected_low_doc_id, "Document text not available.")
+        doc_snippet = full_doc_text[:500] + "..." if len(full_doc_text) > 500 else full_doc_text
+        output_str += f"### Lowly Relevant Document (Qrel <= {low_qrel_threshold})\n"
+        output_str += f"**Document ID:** {selected_low_doc_id}\n"
+        output_str += f"**Document Snippet:** {doc_snippet}\n\n" # Changed to snippet
+    else:
+        output_str += "No lowly relevant document found for this query.\n\n"
+
+    return output_str
+
+
 # --- Initial Load and Indexing Calls ---
 # This part runs once when the app starts.
-# --- IMPORTANT CHANGE: Calling the function that loads Cranfield ---
 load_cranfield_corpus_ir_datasets()
 
 if initial_doc_model_for_indexing == "SPLADE-cocondenser-distil" and model_splade is not None:
@@ -443,7 +576,7 @@ with gr.Blocks(title="SPLADE Demos") as demo:
             )
 
         with gr.TabItem("Document Retrieval Demo"):
-            gr.Markdown("### Retrieve Documents from Cranfield Collection") # Changed title
+            gr.Markdown("### Retrieve Documents from Cranfield Collection")
             gr.Interface(
                 fn=predict_retrieval_gradio,
                 inputs=[
@@ -476,5 +609,15 @@ with gr.Blocks(title="SPLADE Demos") as demo:
                 allow_flagging="never",
                 # live=True # retrieval is too heavy for live
             )
+            
+            gr.Markdown("---") # Separator
+            gr.Markdown("### Get Specific Retrieval Examples")
+            specific_example_output = gr.Markdown()
+            specific_example_button = gr.Button("Get Random Query with High/Low Qrel Docs")
+            specific_example_button.click(
+                fn=get_specific_retrieval_examples,
+                inputs=[],
+                outputs=specific_example_output
+            )
 
-demo.launch()
+demo.launch()