query_documents.py

from pylate import indexes, models, retrieve
import time
import json
from datetime import datetime

# 1. Load model
model = models.ColBERT(model_name_or_path="lightonai/Reason-ModernColBERT")

# 2. Load existing index
index = indexes.Voyager(
    index_folder="./pylate_index",
    index_name="my_docs_index"
)

# 3. Initialize retriever
retriever = retrieve.ColBERT(index=index)

# Load doc_id to filename mapping
with open("doc_id_map.json", "r", encoding="utf-8") as f:
    doc_id_map = json.load(f)

log_file = "search_log.txt"

# 4. Run queries
while True:
    query = input("\nEnter your search query (type 'exit' to quit): ")
    if query.lower() == 'exit':
        break
    
    start_time = time.time()
    
    # Encode query
    query_embedding = model.encode(
        [query],
        is_query=True
    )
    
    # Retrieve results
    results = retriever.retrieve(query_embedding, k=5)

    print("Raw results:", results)  # Debug print
    # Display results
    print(f"\nResults for: '{query}'")
    print(f"Search took: {time.time() - start_time:.2f} seconds")
    print("=" * 50)
    
    # Prepare log entry
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "query": query,
        "results": []
    }
    for hit in results[0]:
        doc_id = hit.get('id', 'N/A')
        score = hit.get('score', 'N/A')
        filename = doc_id_map.get(doc_id, 'Unknown file')
        try:
            score_str = f"{float(score):.4f}"
            print(f"Filename: {filename} | Document ID: {doc_id} | Score: {score_str}")
        except (ValueError, TypeError):
            score_str = str(score)
            print(f"Filename: {filename} | Document ID: {doc_id} | Score: {score_str}")
        print("-" * 50)
        log_entry["results"].append({
            "filename": filename,
            "doc_id": doc_id,
            "score": score_str
        })
    # Write log entry to file
    with open(log_file, "a", encoding="utf-8") as f:
        f.write(json.dumps(log_entry) + "\n")