retriever.py

from sentence_transformers import SentenceTransformer
import numpy as np
import time
from embedder import get_model

# Use the preloaded model from embedder instead of creating a new instance
def retrieve_chunks(index, texts, query, k=5):
    start_time = time.time()
    print(f"Retrieving chunks for query: '{query[:50]}...'")
    
    # Time query embedding
    embed_start = time.time()
    model = get_model()  # Use the preloaded model
    query_vec = model.encode([query])
    embed_time = time.time() - embed_start
    print(f"Query embedding took: {embed_time:.3f} seconds")
    
    # Time FAISS search
    search_start = time.time()
    distances, indices = index.search(np.array(query_vec), k)
    search_time = time.time() - search_start
    print(f"FAISS search took: {search_time:.3f} seconds")
    
    # Time result processing
    process_start = time.time()
    results = [texts[i] for i in indices[0]]
    process_time = time.time() - process_start
    print(f"Result processing took: {process_time:.3f} seconds")
    
    total_time = time.time() - start_time
    print(f"Total chunk retrieval took: {total_time:.3f} seconds")
    print(f"Retrieved {len(results)} chunks")
    
    return results