lambdagent/agentexample/qademo/optimize_index_v2.py

#!/usr/bin/env python3
"""
Optimize RAG index v2: aggressive pruning + numpy arrays
"""
import json
import pickle
import time
import struct
from pathlib import Path

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

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

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

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

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

# More aggressive pruning
MIN_FREQ = 3        # Must appear in >=3 chunks
MAX_FREQ_PCT = 0.15 # Must appear in <15% of chunks
MIN_KW_LEN = 2      # Keyword must be >=2 chars

max_freq = int(len(chunks) * MAX_FREQ_PCT)

pruned = {}
stats = {"rare": 0, "common": 0, "short": 0, "kept": 0}
for kw, posting_list in keywords.items():
    n = len(posting_list)
    if len(kw) < MIN_KW_LEN:
        stats["short"] += 1
        continue
    if n < MIN_FREQ:
        stats["rare"] += 1
        continue
    if n > max_freq:
        stats["common"] += 1
        continue
    # Convert to int indices, deduplicate, sort
    int_postings = sorted(set(id_to_idx[cid] for cid in posting_list if cid in id_to_idx))
    if int_postings:
        # Store as bytes for compact storage (2 bytes per int for <65536 chunks)
        if len(chunks) < 65536:
            pruned[kw] = bytes(struct.pack(f"<{len(int_postings)}H", *int_postings))
        else:
            pruned[kw] = int_postings
    stats["kept"] += 1

print(f"\n--- Pruning stats ---")
for k, v in stats.items():
    print(f"  {k}: {v}")

# Optimized chunks (only essential fields)
opt_chunks = []
for c in chunks:
    opt_chunks.append((c["source"], c["chunk_id"], c["text"]))

# Save
print("\nSaving...")
t0 = time.time()
opt_data = {
    "total": len(opt_chunks),
    "files": data.get("total_files", 0),
    "chunks": opt_chunks,  # list of (source, chunk_id, text) tuples
    "use_bytes": len(chunks) < 65536,
}
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")

idx_size = OPT_INDEX.stat().st_size / 1024 / 1024
kw_size = OPT_KEYWORDS.stat().st_size / 1024 / 1024
print(f"\n--- File sizes ---")
print(f"  Chunks: {idx_size:.1f} MB")
print(f"  Keywords: {kw_size:.1f} MB (was 2439.7 MB JSON)")

# Verify
print("\n--- Load test ---")
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"  Chunks: {(t1-t0)*1000:.0f}ms")
print(f"  Keywords: {(t2-t1)*1000:.0f}ms")
print(f"  Total: {(t2-t0)*1000:.0f}ms (was 50s JSON)")