Adding Baseline and DocumentCollection

Baseline/boolean.py

@@ -0,0 +1,54 @@
+from Baseline.data_processor import process_json_data, process_queries, merge_documents
+from Baseline.boolean_retrieval import main_boolean_retrieval, retrieve_single_query
+import json   
+
+def boolean_pipeline(query, wikipedia_data_path="Datasets/mini_wiki_collection.json", top_n=100):
+    # Load the JSON files
+    with open(wikipedia_data_path, "r") as file1:  
+        wikipedia_data = json.load(file1)
+
+    # Process the JSON files
+    wikipedia_dict = process_json_data(wikipedia_data)
+    # Print the processed data
+
+    top_results = retrieve_single_query(query, wikipedia_dict, top_n)
+
+    return top_results
+
+# def main():
+#     # Load the JSON files
+#     # boolean_retrieval("In the United States, why are positions like Attorney General, Secretary of State, etc. appointed by the president at the federal level but elected by the people at the state level? Had it ever been proposed to do this differently?")
+#     # return
+#     with open("../Datasets/mini_wiki_collection.json", "r") as file1:  # Replace with the actual path to your file
+#         wikipedia_data = json.load(file1)
+
+#     with open("../Datasets/mini_wiki_collection_10000_documents.json", "r") as file1:  # Replace with the actual path to your file
+#         additional_json_file = json.load(file1)
+    
+#     with open("../Datasets/FinalDataset_WithModifiedQuery.json", "r") as file2:  # Replace with the actual path to your file
+#         queries_data = json.load(file2)
+
+#     # Process the JSON files
+#     wikipedia_dict = process_json_data(wikipedia_data)
+#     updated_main_dict = merge_documents(wikipedia_dict, additional_json_file, limit=2000)
+#     queries_dict = process_queries(queries_data)
+
+#     # Print the processed data
+#     print("Processed Wikipedia Data:")
+#     print(wikipedia_dict["420538"])
+#     print("\nProcessed Queries Data:")
+#     print(queries_dict["5xvggq"])
+
+#     top_results = main_boolean_retrieval(updated_main_dict, queries_dict)
+
+#     # Print the results for a specific query
+#     print("\nTop results for query '5xvggq':")
+#     print(top_results.get("5xvggq", []))
+
+#     # Optionally, save the top results to a JSON file
+#     with open("boolean_retrieval_1_2_query.json", "w") as output_file:
+#         json.dump(top_results, output_file, indent=4)
+
+
+# # if __name__ == "__main__":
+# #     main()

Baseline/boolean_retrieval.py

@@ -0,0 +1,134 @@
+from collections import defaultdict
+import re
+import heapq
+import joblib
+import os
+
+def preprocess_text(text):
+    """
+    Preprocess the text for tokenization.
+    Removes special characters, lowercases, and splits into words.
+    """
+    return re.findall(r'\w+', text.lower())
+
+def create_inverted_index(wikipedia_dict):
+    """
+    Create an inverted index from the document dictionary.
+    Args:
+        wikipedia_dict (dict): A dictionary with document IDs as keys and text as values.
+
+    Returns:
+        dict: An inverted index where each term maps to a list of document IDs containing it.
+    """
+    inverted_index = defaultdict(set)
+    for doc_id, text in wikipedia_dict.items():
+        tokens = set(preprocess_text(text))  # Unique tokens for each document
+        for token in tokens:
+            inverted_index[token].add(doc_id)
+    return inverted_index
+
+def save_inverted_index(inverted_index, filepath="Baseline/inverted_index.pkl"):
+    """
+    Save the inverted index to a file using joblib.
+    """
+    joblib.dump(inverted_index, filepath)
+
+def load_inverted_index(filepath="Baseline/inverted_index.pkl"):
+    """
+    Load the inverted index from a file using joblib.
+    """
+    if os.path.exists(filepath):
+        return joblib.load(filepath)
+    return None
+
+def boolean_retrieval(queries_dict, inverted_index, wikipedia_dict, top_n=100):
+    """
+    Perform boolean retrieval for each query.
+    Args:
+        queries_dict (dict): A dictionary with query IDs as keys and query text as values.
+        inverted_index (dict): The inverted index created from the document collection.
+        wikipedia_dict (dict): The original document dictionary (for scoring if needed).
+        top_n (int): The number of top documents to retrieve for each query.
+
+    Returns:
+        dict: A dictionary with query IDs as keys and a list of top document IDs as values.
+    """
+    query_results = {}
+    
+    for query_id, query_text in queries_dict.items():
+        query_tokens = preprocess_text(query_text)
+        
+        # Collect all document IDs that contain any of the query terms
+        relevant_docs = set()
+        for token in query_tokens:
+            if token in inverted_index:
+                relevant_docs.update(inverted_index[token])
+        
+        # If more than `top_n` documents, sort by some criteria (e.g., frequency of terms in the doc)
+        doc_scores = []
+        for doc_id in relevant_docs:
+            doc_text = preprocess_text(wikipedia_dict[doc_id])
+            score = sum(doc_text.count(token) for token in query_tokens)  # Term frequency score
+            doc_scores.append((score, doc_id))
+        
+        # Get the top `top_n` documents based on the score
+        top_docs = heapq.nlargest(top_n, doc_scores)
+        query_results[query_id] = [doc_id for _, doc_id in top_docs]
+
+    return query_results
+
+# Main flow
+def main_boolean_retrieval(wikipedia_dict, queries_dict):
+    # Step 1: Create inverted index
+    inverted_index = create_inverted_index(wikipedia_dict)
+    
+    # Step 2: Perform boolean retrieval
+    top_docs = boolean_retrieval(queries_dict, inverted_index, wikipedia_dict)
+    
+    return top_docs
+
+def retrieve_single_query(query, wikipedia_dict, top_n=100, inverted_index_path="Baseline/inverted_index.pkl"):
+    """
+    Retrieve documents for a single query using the inverted index.
+    If the inverted index is not found, it will be created and saved.
+
+    Args:
+        query (str): The query text.
+        wikipedia_dict (dict): The original document dictionary.
+        top_n (int): The number of top documents to retrieve.
+        inverted_index_path (str): Path to the saved inverted index file.
+
+    Returns:
+        list: A list of top document IDs matching the query.
+    """
+    # Load or create the inverted index
+    inverted_index = load_inverted_index(inverted_index_path)
+    if inverted_index is None:
+        print("Inverted index not found. Creating one...")
+        inverted_index = create_inverted_index(wikipedia_dict)
+        save_inverted_index(inverted_index, inverted_index_path)
+
+    # Preprocess the query
+    query_tokens = preprocess_text(query)
+    
+    # Collect relevant documents
+    relevant_docs = set()
+    for token in query_tokens:
+        if token in inverted_index:
+            relevant_docs.update(inverted_index[token])
+    
+    # Rank documents by frequency of terms
+    doc_scores = []
+    for doc_id in relevant_docs:
+        doc_text = preprocess_text(wikipedia_dict[doc_id])
+        score = sum(doc_text.count(token) for token in query_tokens)
+        doc_scores.append((score, doc_id))
+    
+    # Get the top `top_n` documents based on the score
+    top_docs = heapq.nlargest(top_n, doc_scores)
+    return [doc_id for _, doc_id in top_docs]
+
+# Example usage:
+# Assuming `wikipedia_dict` and `queries_dict` are already prepared
+# top_results = main_boolean_retrieval(wikipedia_dict, queries_dict)
+# print(top_results)

Baseline/data_processor.py

@@ -0,0 +1,98 @@
+# Assuming sanitize_text is a function you've defined elsewhere
+
+import re
+
+def merge_documents(main_dict, additional_json, limit=1000):
+    """
+    Adds a subset of documents from an additional JSON file to the main dictionary.
+
+    Args:
+        main_dict (dict): The main dictionary where processed documents are stored.
+        additional_json (list): The additional JSON data containing documents.
+        limit (int): The maximum number of documents to add to the main dictionary.
+
+    Returns:
+        dict: The updated main dictionary with additional documents added.
+    """
+    # Counter to track how many documents have been added
+    count = 0
+
+    for doc in additional_json:
+        if count >= limit:
+            break
+        
+        # Extract wikipedia_id and text from the document
+        wikipedia_id = doc.get("wikipedia_id")
+        text = doc.get("text", [])
+        
+        # Check if the document ID is unique to avoid overwriting
+        if wikipedia_id not in main_dict:
+            # Process and sanitize the document
+            joined_text = " ".join(text)
+            sanitized_text = sanitize_text(joined_text)
+            
+            # Add to the main dictionary
+            main_dict[wikipedia_id] = sanitized_text
+            count += 1
+    
+    print(f"{count} documents added to the main dictionary.")
+    return main_dict
+
+def sanitize_text(text):
+    """
+    Cleans and standardizes text by keeping only alphanumeric characters and spaces.
+    Args:
+        text (str): Text to sanitize.
+    Returns:
+        str: Sanitized text.
+    """
+    if isinstance(text, str):
+        # Use regex to keep only alphanumeric characters and spaces
+        text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
+        # Optionally, collapse multiple spaces into a single space
+        text = re.sub(r'\s+', ' ', text).strip()
+    return text
+
+
+def process_json_data(json_data):
+    result_dict = {}
+    
+    for doc in json_data:
+        # Extract wikipedia_id and text
+        wikipedia_id = doc.get("wikipedia_id")
+        text = doc.get("text", [])
+        
+        # Join the text content and sanitize
+        joined_text = " ".join(text)
+        sanitized_text = sanitize_text(joined_text)
+        
+        # Store in the dictionary
+        result_dict[wikipedia_id] = sanitized_text
+
+    return result_dict
+
+def process_queries(json_data):
+    """
+    Processes a JSON object containing queries and query IDs.
+
+    Args:
+        json_data (dict): The input JSON data.
+
+    Returns:
+        dict: A dictionary with query_id as the key and query text as the value.
+    """
+    result_dict = {}
+    
+    for query_id, query_info in json_data.items():
+        # Extract the query input
+        query_text = query_info.get("input", "")
+        
+        # Store query_id and text in the result dictionary
+        result_dict[query_id] = query_text
+
+    return result_dict
+
+# Example usage
+# Assuming `query_json_file` contains your JSON data
+# processed_queries = process_queries(query_json_file)
+

Baseline/inverted_index.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c47f19521041e7b2a5681da4128cfae538eba1bc653528f04c7dc9df300fbc5
+size 4671080

Datasets/mini_wiki_collection.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:986eedb174550564ce95cf9b08de1207cfb1e2290646b4aeb60257c9edceb27a
+size 41656963