File size: 2,758 Bytes
6a02b16 c6cece9 6a02b16 | 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 | import json
from pathlib import Path
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from data_io import load_clauses, write_jsonl
def build_one(model_name, out_dir, ru, kz, batch_size=64):
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
ru_texts = [x["text"] for x in ru]
kz_texts = [x["text"] for x in kz]
ru_meta = [{"pos": i, "id": ru[i]["id"], "lang": "ru", "meta": ru[i]["meta"], "text": ru[i]["text"]} for i in range(len(ru))]
kz_meta = [{"pos": i, "id": kz[i]["id"], "lang": "kz", "meta": kz[i]["meta"], "text": kz[i]["text"]} for i in range(len(kz))]
model = SentenceTransformer(model_name)
ru_emb = model.encode(ru_texts, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True).astype(np.float32)
kz_emb = model.encode(kz_texts, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True).astype(np.float32)
ru_index = faiss.IndexFlatIP(ru_emb.shape[1])
ru_index.add(ru_emb)
faiss.write_index(ru_index, str(out_dir / "ru.faiss"))
kz_index = faiss.IndexFlatIP(kz_emb.shape[1])
kz_index.add(kz_emb)
faiss.write_index(kz_index, str(out_dir / "kz.faiss"))
write_jsonl(str(out_dir / "ru_meta.jsonl"), ru_meta)
write_jsonl(str(out_dir / "kz_meta.jsonl"), kz_meta)
info = {
"model": model_name,
"ru_count": len(ru_texts),
"kz_count": len(kz_texts),
"ru_index": str(out_dir / "ru.faiss"),
"kz_index": str(out_dir / "kz.faiss"),
"ru_meta": str(out_dir / "ru_meta.jsonl"),
"kz_meta": str(out_dir / "kz_meta.jsonl"),
}
(out_dir / "index_info.json").write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")
def main():
clauses_path = "data/clauses_constitution_ru_kz.jsonl"
ru, kz = load_clauses(clauses_path)
base_root = Path("artifacts/indexes")
base_root.mkdir(parents=True, exist_ok=True)
models = [
("mpnet_base", "paraphrase-multilingual-mpnet-base-v2"),
("labse", "sentence-transformers/LaBSE"),
]
finetuned_path = Path("artifacts/models/finetuned_mpnet")
if finetuned_path.exists():
models.append(("finetuned", str(finetuned_path)))
built = []
for alias, model_name in models:
out_dir = base_root / alias
build_one(model_name, out_dir, ru, kz, batch_size=64)
built.append({"alias": alias, "dir": str(out_dir), "model": model_name})
(base_root / "built_indexes.json").write_text(json.dumps(built, ensure_ascii=False, indent=2), encoding="utf-8")
if __name__ == "__main__":
main()
|