File size: 6,397 Bytes
83892b0 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """
indexer.py
──────────
Reads articles.jsonl, embeds each article using a multilingual
sentence-transformers model, and saves a FAISS index to disk.
What this script produces:
data/faiss.index — the vector index (for similarity search)
data/metadata.jsonl — article metadata in the same order as the index
(needed to recover the text after a search)
Usage:
python indexer.py
Install:
pip install faiss-cpu sentence-transformers
"""
import json
import faiss
import numpy as np
from pathlib import Path
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
# ── Config ────────────────────────────────────────────────────────────────────
MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
ARTICLES_FILE = Path("data/articles.jsonl")
INDEX_FILE = Path("data/faiss.index")
METADATA_FILE = Path("data/metadata.jsonl")
BATCH_SIZE = 32 # embed this many articles at a time
# ── Load articles ─────────────────────────────────────────────────────────────
def load_articles(path: Path) -> list[dict]:
articles = []
with open(path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
articles.append(json.loads(line))
print(f"Loaded {len(articles)} articles from {path}")
return articles
# ── Embed ─────────────────────────────────────────────────────────────────────
def embed_articles(articles: list[dict], model: SentenceTransformer) -> np.ndarray:
"""
Convert each article's 'chunk' field into a vector.
The 'chunk' field contains:
"LEGE 53 2003\nArticolul 142\n\nSalariații au dreptul..."
Including the law title and article number in the chunk means the
embedding captures source context, not just the article text alone.
This improves retrieval accuracy.
Returns a 2D numpy array of shape (n_articles, embedding_dim).
"""
chunks = [a["chunk"] for a in articles]
all_embeddings = []
print(f"Embedding {len(chunks)} articles in batches of {BATCH_SIZE}...")
for i in tqdm(range(0, len(chunks), BATCH_SIZE)):
batch = chunks[i : i + BATCH_SIZE]
embeddings = model.encode(
batch,
convert_to_numpy=True,
normalize_embeddings=True, # normalize for cosine similarity
show_progress_bar=False,
)
all_embeddings.append(embeddings)
return np.vstack(all_embeddings).astype("float32")
# ── Build FAISS index ─────────────────────────────────────────────────────────
def build_index(embeddings: np.ndarray) -> faiss.Index:
"""
Build a FAISS index from the embeddings.
We use IndexFlatIP (Inner Product) because our embeddings are
normalized — inner product == cosine similarity when vectors are
unit length. This gives us semantic similarity search for free.
For larger datasets (100k+ articles) you'd switch to IndexIVFFlat
which is approximate but much faster. For our size, exact search is fine.
"""
dim = embeddings.shape[1] # embedding dimension (384 for MiniLM)
print(f"Building FAISS index — {len(embeddings)} vectors, dim={dim}")
index = faiss.IndexFlatIP(dim) # IP = Inner Product (cosine similarity)
index.add(embeddings)
print(f"Index contains {index.ntotal} vectors")
return index
# ── Save ──────────────────────────────────────────────────────────────────────
def save_metadata(articles: list[dict], path: Path):
"""
Save article metadata in the same order as the FAISS index.
FAISS only stores vectors — it doesn't store text. When we search
and get back index positions [42, 17, 8], we need this file to
look up what article #42 actually says.
"""
with open(path, "w", encoding="utf-8") as f:
for article in articles:
# Save only what we need for display — not the full chunk
meta = {
"law_id": article.get("law_id"),
"law_title": article.get("law_title"),
"article_number": article.get("article_number"),
"text": article.get("text"),
"chunk": article.get("chunk"),
}
f.write(json.dumps(meta, ensure_ascii=False) + "\n")
print(f"Saved metadata to {path}")
# ── Main ──────────────────────────────────────────────────────────────────────
def main():
# 1. Load articles
articles = load_articles(ARTICLES_FILE)
if not articles:
print("No articles found. Run pipeline.py first.")
return
# 2. Load embedding model
# First run downloads ~120MB — subsequent runs use the cache
print(f"Loading model: {MODEL_NAME}")
model = SentenceTransformer(MODEL_NAME)
# 3. Embed all articles
embeddings = embed_articles(articles, model)
# 4. Build FAISS index
index = build_index(embeddings)
# 5. Save index and metadata
faiss.write_index(index, str(INDEX_FILE))
print(f"Saved FAISS index to {INDEX_FILE}")
save_metadata(articles, METADATA_FILE)
print("\nDone! You can now run retriever.py to search the index.")
print(f" Index: {INDEX_FILE} ({INDEX_FILE.stat().st_size // 1024} KB)")
print(f" Metadata: {METADATA_FILE} ({METADATA_FILE.stat().st_size // 1024} KB)")
if __name__ == "__main__":
main() |