evaluation_pipeline/retrieval.py

import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np
import tracemalloc
import time

import pandas as pd
import numpy as np
import sqlite3
import sqlite_vec
from typing import List
import struct
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import seaborn as sns
import logging
import os
import sys
import argparse
import pickle

# Add the parent directory of `evaluation_pipeline` to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from src.constants import EMBEDDING_MODELS_DICT
from src.feature_extractor import FeatureExtractor


def log_performance(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        tracemalloc.start()
        
        result = func(*args, **kwargs)
        
        current, peak = tracemalloc.get_traced_memory()
        end_time = time.time()
        
        logging.info(f"Function: {func.__name__}")
        logging.info(f"Execution Time: {end_time - start_time:.2f} seconds")
        logging.info(f"Current Memory Usage: {current / 10**6:.2f} MB")
        logging.info(f"Peak Memory Usage: {peak / 10**6:.2f} MB")
        
        tracemalloc.stop()
        return result
    return wrapper

def format_size(size_in_bytes):
    """Convert bytes to a human-readable format (KB, MB, GB, etc.)."""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_in_bytes < 1024:
            return f"{size_in_bytes:.2f} {unit}"
        size_in_bytes /= 1024


def process_golden_queries(golden_query_file_path):
    golden_df = pd.read_csv(golden_query_file_path)
    query_ids = {query: str(hash(query)) for query in golden_df['search_query'].unique()}
    return query_ids


@log_performance
def process_history(row_limit, history_file_path):
    browsing_history = pd.read_csv(history_file_path).head(row_limit)
    browsing_history['last_visit_date'] = pd.to_datetime(browsing_history['last_visit_date'], unit='us')
    # fill empty last_visit_date with default value "1970-01-01"
    browsing_history['last_visit_date'] = browsing_history['last_visit_date'].fillna(pd.to_datetime("1970-01-01"))
    browsing_history['combined_text'] = browsing_history['title'].fillna('') + " " + browsing_history['description'].fillna('')
    browsing_history['combined_text_url'] = browsing_history['title'].fillna('') + " " + browsing_history['description'].fillna('') + browsing_history['url'].fillna('')
    browsing_history = browsing_history.loc[browsing_history['combined_text'] != ''].reset_index(drop=True)

    print(len(browsing_history))

    return browsing_history

@log_performance
def create_embeddings(row_limit, history, embeddings_model_dict):
    texts = history['combined_text'].values.tolist()
    embeddings_dict = {}
    embeddings_sizes = {}

    for model in embeddings_model_dict.keys():
        if model == 'nomic-ai/nomic-embed-text-v1.5':
            prefix = 'search_document: '
            texts = [prefix + text for text in texts]
        fe = FeatureExtractor(embeddings_model_dict, model_name=model)
        embeddings_dict[model] = fe.get_embeddings(texts)
        print(model, embeddings_dict[model].shape)
        embeddings_sizes[model] = embeddings_dict[model].shape[1]

    with open(f"data/embeddings_dict_{row_limit}.pkl", "wb") as f:
        pickle.dump(embeddings_dict, f)

    with open(f"data/embeddings_sizes_{row_limit}.pkl", "wb") as f:
        pickle.dump(embeddings_sizes, f)

    return embeddings_dict, embeddings_sizes

@log_performance
def create_db():
    # vector database
    #db = sqlite3.connect(":memory:")
    db = sqlite3.connect(":memory:")
    db.enable_load_extension(True)
    sqlite_vec.load(db)
    db.enable_load_extension(False)

    sqlite_version, vec_version = db.execute(
        "select sqlite_version(), vec_version()"
    ).fetchone()
    print(f"sqlite_version={sqlite_version}, vec_version={vec_version}")
    return db

def serialize_f32(vector: List[float]) -> bytes:
    """serializes a list of floats into a compact "raw bytes" format"""
    return struct.pack("%sf" % len(vector), *vector)

@log_performance
def create_embeddings_table_in_vector_db(db, model_name, embeddings_sizes, embeddings_dict):
    print("creating table")
    EMBEDDING_SIZE = embeddings_sizes[model_name]
    items = []
    for idx, vec in enumerate(embeddings_dict[model_name]):
        items.append((idx, list(vec)))
    model_name_normalized = model_name.replace("/","_").replace("-","_").replace(".","_")
    db.execute(f"CREATE VIRTUAL TABLE vec_items_{model_name_normalized} USING vec0(embedding float[{EMBEDDING_SIZE}])")

    with db:
        for item in items:
            db.execute(
                f"INSERT INTO vec_items_{model_name_normalized}(rowid, embedding) VALUES (?, ?)",
                [item[0], serialize_f32(item[1])],
            )
    return db


# retrieval
@log_performance
def query_and_result(fe, query, db, model_name, threshold, k):
    model_name_normalized = model_name.replace("/","_").replace("-","_").replace(".","_") 
    if model_name == 'nomic-ai/nomic-embed-text-v1.5':
        query = 'search_query: ' + query
   
    query_embedding = fe.get_embeddings([query])[0]
    # using cosine distance
    rows = db.execute(
    f"""
      SELECT
        a.rowid,
        vec_distance_cosine(embedding, ?) AS cosine_distance
      FROM vec_items_{model_name_normalized} a
      inner join search_data b
      on a.rowid = b. rowid

      where b.url not like '%google.com/search?%'
      ORDER BY cosine_distance
      LIMIT {k}
    """,
    [serialize_f32(query_embedding)],
    ).fetchall()

    results = []

    for row in rows:
        rowid = row[0]  # Get the row ID
        distance = row[1]  # Get the cosine distance

        # Skip rows where distance > threshold
        if distance > threshold:
            print("Distance threshold exceeded")
            continue

        print("doing search now")
        # Step 2: Query additional details for the matching row from search_data
        res = db.execute(
            """
            SELECT rowid, title, url, combined_text
            FROM search_data
            WHERE rowid = ?
            """,
            (rowid,)
        ).fetchone()

        # Add the result to the results list
        if res:
            results.append({
                "id": res[0],
                "title": res[1],
                "url": res[2],
                "combined_text": res[3],
                "distance": distance,
            })

    return results

def get_table_size(connection, table_name):
    """Get the approximate size of a table in SQLite."""
    with connection:
        cursor = connection.cursor()
        # Get page size and page count
        cursor.execute("PRAGMA page_count;")
        page_count = cursor.fetchone()[0]

        cursor.execute("PRAGMA page_size;")
        page_size = cursor.fetchone()[0]
        
        total_db_size = page_count * page_size

        # If dbstat is available, use it for more precision
        try:
            cursor.execute(f"""
                SELECT SUM(pgsize) AS table_size
                FROM dbstat
                WHERE name = ?;
            """, (table_name,))
            result = cursor.fetchone()
            if result and result[0]:
                return result[0]
        except sqlite3.DatabaseError:
            print("dbstat is not available; estimating based on database size.")

        # Fallback to database size as an estimate
        return total_db_size



@log_performance
def load_history_in_db(db, browsing_history):
    db.execute('''CREATE TABLE IF NOT EXISTS search_data (
        url TEXT,
        title TEXT,
        description TEXT,
        combined_text TEXT
        )
        ''')
    for idx, row in browsing_history.iterrows():
            db.execute("""
                INSERT INTO search_data (rowid, url, title, description, combined_text)
                VALUES (?, ?, ?, ?, ?)
            """, (idx, row['url'], row['title'], row['description'],  row['combined_text'])
            )
    browsing_history.to_sql("full_history", db, if_exists="replace", index=True)  # Creates the table automatically



def create_ground_truth(history, query_ids):
    queries = []
    relevant_ids = []
    for query in history['search_query'].unique():
        query_id = query_ids[query]
        queries.append(query_id)
        ground_truths = []
        for i, row in history.iterrows():
            relevant = 1 if query == row['search_query'] else 0
            if relevant == 1:
                ground_truths.append(int(i))
        relevant_ids.append(ground_truths)

    return dict(zip(queries, relevant_ids))

def load_ground_truth_from_golden(db, golden_df_file_path):
    golden_df = pd.read_csv(golden_df_file_path)
    query_ids = {query: str(hash(query)) for query in golden_df['search_query'].unique()}
    db.execute('''CREATE TABLE IF NOT EXISTS ground_truth (
        search_query TEXT,
        url TEXT
        )
        ''')
    with db:
       for _, row in golden_df.iterrows():
            db.execute("""
                INSERT INTO ground_truth (search_query, url)
                VALUES (?, ?)
            """, ( row['search_query'], row['url'])
            )

    # join to search query to get doc ID for ground truth URL
    results = db.execute(
        '''SELECT a.search_query, a.url, b.rowid
        FROM ground_truth a
        left join search_data b
        on a.url = b.url'''
    ).fetchall()

    ground_truth = {}
    ground_truth_urls = {}
    for query, url, id_ in results:
        query_id = query_ids[query]
        if query_id not in ground_truth:
            ground_truth[query_id] = []
            ground_truth_urls[query_id] = []
            
        ground_truth[query_id].append(id_)
        ground_truth_urls[query_id].append(url)

    return ground_truth, query_ids, ground_truth_urls


@log_performance
def run_history_in_vector_db(row_limit, history_file_path, golden_set_file_path):

    browsing_history = process_history(row_limit, history_file_path=history_file_path)

    # create vector DB
    db = create_db()


    # load in history for joining later
    load_history_in_db(db, browsing_history)


    # if a golden set is not provided, assume it's with the history
    if not golden_set_file_path:
        query_ids = {query: str(hash(query)) for query in browsing_history['search_query'].unique()}
        browsing_history['query_id'] = browsing_history['search_query'].map(query_ids)
        ground_truth = create_ground_truth(browsing_history, query_ids)
    else:
        print("Getting doc ids for history")
        ground_truth, query_ids, ground_truth_urls = load_ground_truth_from_golden(db, golden_df_file_path=golden_set_file_path)



    # create embeddings for candidate models
    print("Generating Embeddings")
    try:
        path = f"data/embeddings_dict_{row_limit}.pkl"
        # path = f"/Users/rebeccahadi/Documents/search-your-history-poc/data/embeddings_dict_{row_limit}.pkl"
        with open(path, "rb") as f:
            embeddings_dict = pickle.load(f)

        sizes_path = f"data/embeddings_sizes_{row_limit}.pkl"
     #   sizes_path = f"/Users/rebeccahadi/Documents/search-your-history-poc/data/embeddings_sizes_{row_limit}.pkl"
        with open(sizes_path, "rb") as f:
            embeddings_sizes = pickle.load(f)
    except:
        embeddings_dict, embeddings_sizes = create_embeddings(row_limit, browsing_history, embeddings_model_dict=EMBEDDING_MODELS_DICT)

    # loop through each model/embedding type and store in db
    for model_name in embeddings_dict.keys():
        model_name_normalized = model_name.replace("/","_").replace("-","_").replace(".","_")

        # create table for embeddings for model
        create_embeddings_table_in_vector_db(db, model_name, embeddings_sizes=embeddings_sizes, embeddings_dict=embeddings_dict)

        table_size = get_table_size(db, table_name=model_name_normalized)
        logging.info(f"{model_name_normalized} table size: {table_size}")
        total_db_size_human_readable = format_size(table_size)
        logging.info(f"Table size {model_name_normalized}: {total_db_size_human_readable}")

    return query_ids, db, ground_truth, ground_truth_urls

@log_performance
def run_retrieval(fe, query_ids, db, model_name, threshold, k):
    # loop through each query
    retreival_dict = {}
    query_lookup = {}
    for query in query_ids.keys():
        # for later identifying query
        query_id = query_ids[query]
        query_lookup[query_id] = query
        # perform retrieval
        results = query_and_result(fe, query, db=db,
        model_name=model_name, threshold=threshold, k=k)
        retreival_dict[query_id] = results
    return retreival_dict, query_lookup


def convert_dict_to_df(retrieval_dict, query_lookup, ground_truth, ground_truth_urls, model_name, k):
    rows = []
    for query_id, retrievals in retrieval_dict.items():
        # Flatten each retrieval into a single row with column names based on retrieval index
        row = {'query_id': str(query_id)}
        retrieved_ids = []  # List to collect all retrieved IDs
        retrieved_distances = [] # collect the distances
        for i, retrieval in enumerate(retrievals, start=1):
            row[f'retrieval_{i}_id'] = retrieval.get('id')
            row[f'retrieval_{i}_title'] = retrieval.get('title')
            row[f'retrieval_{i}_url'] = retrieval.get('url')
            row[f'retrieval_{i}_combined_text'] = retrieval.get('combined_text')
            row[f'retrieval_{i}_distance'] = retrieval.get('distance')
            retrieved_ids.append(retrieval.get('id'))
            retrieved_distances.append(retrieval.get('distance'))
            # Collect the ID for the list

        row['retrieved_ids'] = retrieved_ids
        row['retrieved_distances'] = retrieved_distances
        row['model_name'] = model_name
        row['query'] = query_lookup[query_id]
        row['relevant_docs'] = ground_truth[query_id]
        row['relevant_urls'] = ground_truth_urls[query_id]
        row['k'] = k
        rows.append(row)
    df = pd.DataFrame(rows)
    return df


def main(model_name, k, threshold, history_file_path, golden_path=None, row_limit=100):
    model_name_normalized = model_name.replace("/","_").replace("-","_").replace(".","_")
    # Configure logging
    logging.basicConfig(
    filename=f"performance_{model_name_normalized}.log",
    level=logging.INFO,
    format="%(asctime)s - %(message)s"
    )
    query_ids, db, ground_truth, ground_truth_urls = run_history_in_vector_db(row_limit, history_file_path=history_file_path, golden_set_file_path=golden_path)
    fe = FeatureExtractor(EMBEDDING_MODELS_DICT, model_name=model_name)
    retrieval_results, query_lookup = run_retrieval(fe, query_ids, db, model_name, threshold, k)
    # reshape & save to df and csv
    df = convert_dict_to_df(retrieval_dict=retrieval_results, query_lookup=query_lookup, ground_truth=ground_truth, ground_truth_urls=ground_truth_urls, model_name=model_name, k=k)
    time_stamp = int(time.time())
    df.to_csv(f"results/{model_name_normalized}_results.csv", index=False)
    return db, retrieval_results, df


if __name__ == "__main__":
     # Create the argument parser
     parser = argparse.ArgumentParser(description="Run the retrieval pipeline with specified parameters.")
      # Add arguments
     parser.add_argument("history_file_path", type=str, help="Path to the browsing history file.")
     parser.add_argument("--model_name", type=str,default='Xenova/all-MiniLM-L6-v2', help="Name of the model to use.")
     parser.add_argument("--k", type=int, default=2, help="Top-K results to retrieve.")
     parser.add_argument("--threshold", type=float, default=10.0, help="Threshold for retrieval.")
     parser.add_argument("--golden_path", type=str, default=None, help="Path to the golden query set file (optional).")
     parser.add_argument("--row_limit", type=int, default=100, help="Whether to limit rows from browsing history upon load")
     # Parse arguments
     args = parser.parse_args()
     # Call the main function with parsed arguments
     main(
         model_name=args.model_name,
         k=args.k,
         threshold=args.threshold,
         history_file_path=args.history_file_path,
         golden_path=args.golden_path,
         row_limit=args.row_limit
     )