lambdagent/agentexample/qademo/optimize_index.py

#!/usr/bin/env python3
"""
Optimize RAG keyword index:
1. Prune useless keywords (too frequent / too rare)
2. Use pickle for fast loading
3. Store chunk_id as int index instead of string
"""
import json
import pickle
import time
import re
from pathlib import Path
from collections import Counter

BASE = Path("/data/knowledge/maritime")
INDEX_FILE = BASE / "rag_index.json"
KEYWORD_FILE = BASE / "rag_index.keywords.json"
# Optimized outputs
OPT_INDEX = BASE / "rag_index.pkl"
OPT_KEYWORDS = BASE / "rag_keywords.pkl"

print("=== RAG Index Optimizer ===")

# Load original
print("Loading chunk index...")
t0 = time.time()
with open(INDEX_FILE) as f:
    data = json.load(f)
chunks = data["chunks"]
print(f"  Chunks: {len(chunks)}, loaded in {time.time()-t0:.1f}s")

print("Loading keyword index...")
t0 = time.time()
with open(KEYWORD_FILE) as f:
    keywords = json.load(f)
t1 = time.time()
print(f"  Keywords: {len(keywords)}, loaded in {t1-t0:.1f}s")

# Analyze
lengths = [len(v) for v in keywords.values()]
print(f"\n--- Before optimization ---")
print(f"  Total keywords: {len(keywords)}")
print(f"  Total postings: {sum(lengths)}")
print(f"  Avg postings/keyword: {sum(lengths)/len(lengths):.1f}")
print(f"  Max postings: {max(lengths)}")

# Build chunk_id -> int index mapping
chunk_ids = [c["id"] for c in chunks]
id_to_idx = {cid: i for i, cid in enumerate(chunk_ids)}

# Prune keywords
MIN_FREQ = 2       # Keywords appearing in <2 chunks are noise
MAX_FREQ_PCT = 0.3  # Keywords appearing in >30% of chunks are too common (stopwords)
max_freq = int(len(chunks) * MAX_FREQ_PCT)

pruned = {}
removed_rare = 0
removed_common = 0
for kw, posting_list in keywords.items():
    n = len(posting_list)
    if n < MIN_FREQ:
        removed_rare += 1
        continue
    if n > max_freq:
        removed_common += 1
        continue
    # Convert string IDs to int indices for compact storage
    int_postings = []
    for cid in posting_list:
        if cid in id_to_idx:
            int_postings.append(id_to_idx[cid])
    if int_postings:
        pruned[kw] = int_postings

print(f"\n--- After pruning ---")
print(f"  Removed rare (<{MIN_FREQ}): {removed_rare}")
print(f"  Removed common (>{max_freq}): {removed_common}")
print(f"  Remaining keywords: {len(pruned)}")
print(f"  Total postings: {sum(len(v) for v in pruned.values())}")

# Build optimized chunk data (minimal fields)
opt_chunks = []
for c in chunks:
    opt_chunks.append({
        "id": c["id"],
        "source": c["source"],
        "chunk_id": c["chunk_id"],
        "text": c["text"],
        "length": c["length"],
    })

# Save as pickle
print("\nSaving optimized index (pickle)...")
t0 = time.time()
opt_data = {
    "total_chunks": len(opt_chunks),
    "total_files": data.get("total_files", 0),
    "chunk_size": data.get("chunk_size", 512),
    "chunks": opt_chunks,
    "chunk_ids": chunk_ids,  # for reverse lookup
}
with open(OPT_INDEX, "wb") as f:
    pickle.dump(opt_data, f, protocol=pickle.HIGHEST_PROTOCOL)

with open(OPT_KEYWORDS, "wb") as f:
    pickle.dump(pruned, f, protocol=pickle.HIGHEST_PROTOCOL)

t1 = time.time()
print(f"  Saved in {t1-t0:.1f}s")

# File sizes
idx_size = OPT_INDEX.stat().st_size / 1024 / 1024
kw_size = OPT_KEYWORDS.stat().st_size / 1024 / 1024
old_kw_size = KEYWORD_FILE.stat().st_size / 1024 / 1024
print(f"\n--- File sizes ---")
print(f"  Chunk index: {idx_size:.1f} MB (pickle)")
print(f"  Keyword index: {kw_size:.1f} MB (pickle) vs {old_kw_size:.1f} MB (JSON) = {kw_size/old_kw_size*100:.0f}%")

# Verify load time
print("\n--- Load time verification ---")
t0 = time.time()
with open(OPT_INDEX, "rb") as f:
    d1 = pickle.load(f)
t1 = time.time()
with open(OPT_KEYWORDS, "rb") as f:
    d2 = pickle.load(f)
t2 = time.time()
print(f"  Chunk index load: {(t1-t0)*1000:.0f}ms")
print(f"  Keyword index load: {(t2-t1)*1000:.0f}ms")
print(f"  Total: {(t2-t0)*1000:.0f}ms")
print(f"\n  Speedup: {43000/(t2-t0)/1000:.0f}x faster!")