File size: 2,850 Bytes
187a491 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """Converte metadata.pkl (5GB, caricato interamente in RAM) in SQLite su disco.
Il RAG cerca in FAISS che restituisce indici di riga (0..N-1). Con SQLite
mappiamo direttamente idx -> (title, text) tramite PRIMARY KEY: lookup O(1)
su disco, apertura istantanea, ZERO RAM caricata al cold start.
metadata.pkl contiene:
chunks: dict chunk_id -> {title, text}
chunk_ids: list faiss_idx -> chunk_id
Uso:
python scripts/convert_metadata_to_sqlite.py \
--input rag/index/faiss/metadata.pkl \
--output rag/index/faiss/metadata.db
"""
from __future__ import annotations
import argparse
import os
import pickle
import sqlite3
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="Convert metadata.pkl to SQLite.")
p.add_argument("--input", default="rag/index/faiss/metadata.pkl")
p.add_argument("--output", default="rag/index/faiss/metadata.db")
return p.parse_args()
def main() -> None:
args = parse_args()
print(f"Caricamento {args.input} (puo' richiedere 1-2 min)...")
with open(args.input, "rb") as f:
meta = pickle.load(f)
chunks = meta["chunks"]
chunk_ids = meta.get("chunk_ids", list(chunks.keys()))
total = meta.get("total_chunks", len(chunk_ids))
embed_model = meta.get("embedding_model", "intfloat/multilingual-e5-small")
embed_dim = meta.get("embed_dim", 384)
print(f" {total:,} chunk, modello={embed_model}, dim={embed_dim}")
if os.path.exists(args.output):
os.remove(args.output)
conn = sqlite3.connect(args.output)
cur = conn.cursor()
# PRAGMA per scrittura veloce
cur.execute("PRAGMA journal_mode = OFF")
cur.execute("PRAGMA synchronous = OFF")
cur.execute("""
CREATE TABLE chunks (
idx INTEGER PRIMARY KEY,
title TEXT,
text TEXT
)
""")
# Tabella info (metadata dell'indice)
cur.execute("CREATE TABLE info (key TEXT PRIMARY KEY, value TEXT)")
cur.executemany("INSERT INTO info VALUES (?, ?)", [
("total_chunks", str(total)),
("embedding_model", embed_model),
("embed_dim", str(embed_dim)),
])
print("Scrittura righe in SQLite...")
batch = []
BATCH = 50_000
for idx, cid in enumerate(chunk_ids):
c = chunks[cid]
batch.append((idx, c["title"], c["text"]))
if len(batch) >= BATCH:
cur.executemany("INSERT INTO chunks VALUES (?, ?, ?)", batch)
batch.clear()
if (idx + 1) % 500_000 == 0:
print(f" {idx + 1:,}/{total:,}...")
if batch:
cur.executemany("INSERT INTO chunks VALUES (?, ?, ?)", batch)
conn.commit()
conn.close()
size_gb = os.path.getsize(args.output) / 1e9
print(f"Fatto: {args.output} ({size_gb:.2f} GB)")
if __name__ == "__main__":
main()
|