| |
| """ |
| build_vector_db.py – SOTA Hybrid Search Engine & Dual-Format Artifact Builder |
| |
| Generates: |
| 1. `dpdp_hybrid_index.pkl` -> Python-native serialized hybrid index (Backward compatibility). |
| 2. `dpdp_embeddings.safetensors` -> Zero-copy float32 dense embeddings matrix for Rust Native Daemon. |
| 3. `dpdp_index.json` -> Structured chunks, legal metadata, and Okapi BM25 statistics for Rust Native Daemon. |
| """ |
|
|
| import os |
| import re |
| import sys |
| import json |
| import datetime |
| from pathlib import Path |
| from collections import Counter |
| from typing import List, Dict, Any, Tuple |
|
|
| import numpy as np |
| import pickle |
|
|
| from rank_bm25 import BM25Okapi |
| from sentence_transformers import SentenceTransformer |
| from huggingface_hub import snapshot_download |
| from langchain_text_splitters import RecursiveCharacterTextSplitter |
|
|
| try: |
| from safetensors.numpy import save_file, load_file |
| HAS_SAFETENSORS = True |
| except ImportError: |
| HAS_SAFETENSORS = False |
|
|
| if sys.stdout.encoding != 'utf-8': |
| sys.stdout.reconfigure(encoding='utf-8') |
|
|
| |
| |
| |
| _CURRENT_DIR = Path(__file__).resolve().parent |
| _EVALS_DIR = _CURRENT_DIR.parent / "evals" |
|
|
| if str(_EVALS_DIR) not in sys.path: |
| sys.path.insert(0, str(_EVALS_DIR)) |
|
|
| try: |
| from path_resolver import Paths |
| LAW_TEXT_PATH = Paths.LAW_TEXT |
| HYBRID_INDEX_PATH = Paths.HYBRID_INDEX |
| except ImportError: |
| LAW_TEXT_PATH = _CURRENT_DIR / "dpdp_act_and_rules_2025.txt" |
| HYBRID_INDEX_PATH = _CURRENT_DIR / "dpdp_hybrid_index.pkl" |
|
|
| RAG_DIR = _CURRENT_DIR.parent / "models" / "rag-index" |
| SAFETENSORS_PATH = RAG_DIR / "dpdp_embeddings.safetensors" |
| INDEX_JSON_PATH = RAG_DIR / "dpdp_index.json" |
|
|
|
|
| |
| |
| |
| GENERIC_STOPWORDS = { |
| "the", "a", "an", "is", "are", "was", "were", "of", "and", "in", |
| "to", "for", "with", "on", "at", "by", "from", "as", "that", "this", |
| "it", "be", "or", "which", "will", "would", "could", "should", "their", "they" |
| } |
|
|
| def get_standard_tokenizer(): |
| def tokenize(text: str) -> List[str]: |
| words = re.findall(r'\w+', str(text).lower()) |
| return [w for w in words if w not in GENERIC_STOPWORDS] |
| return tokenize |
|
|
|
|
| |
| |
| |
| def extract_metadata(header_text: str, body_text: str) -> Dict[str, str]: |
| meta = { |
| "type": "General", |
| "number": "", |
| "sub_section": "", |
| "parent_act": "DPDP Act 2023", |
| "applies_to": "all" |
| } |
| |
| full_text = header_text + "\n" + body_text |
| |
| if "RULES 2025" in full_text.upper(): |
| meta["parent_act"] = "DPDP Rules 2025" |
| |
| sec_match = re.search(r'\bSection\s+\d+(?:\.\d+)*', header_text, re.IGNORECASE) |
| rule_match = re.search(r'\bRule\s+\d+(?:\.\d+)*', header_text, re.IGNORECASE) |
| sched_match = re.search(r'\b(?:FIRST|SECOND|THIRD|FOURTH|FIFTH|SIXTH|SEVENTH|EIGHTH)\s+SCHEDULE\b', header_text, re.IGNORECASE) |
| ch_match = re.search(r'\bCHAPTER\s+([IVXLCDM]+)', header_text, re.IGNORECASE) |
| |
| if sec_match: |
| meta["type"] = "Section" |
| meta["number"] = sec_match.group(0) |
| elif rule_match: |
| meta["type"] = "Rule" |
| meta["number"] = rule_match.group(0) |
| elif sched_match: |
| meta["type"] = "Schedule" |
| meta["number"] = sched_match.group(0).upper() |
| elif ch_match: |
| meta["type"] = "Chapter" |
| meta["number"] = ch_match.group(1) |
| |
| sub_match = re.search(r'^\s*\((\d+[a-zA-Z]*)\)', body_text) |
| if sub_match: |
| meta["sub_section"] = f"({sub_match.group(1)})" |
| |
| if "SECOND SCHEDULE" in header_text.upper() or "SEVENTH SCHEDULE" in header_text.upper(): |
| meta["applies_to"] = "state" |
| |
| return meta |
|
|
|
|
| |
| |
| |
| def parse_and_chunk(file_path: Path) -> Tuple[List[str], List[Dict[str, str]]]: |
| with open(file_path, "r", encoding="utf-8") as f: |
| text = f.read() |
| |
| chunks = [] |
| metadatas = [] |
| |
| split_pattern = r'(?=\n+(?:Section\s+\d+(?:\.\d+)*|Rule\s+\d+(?:\.\d+)*|CHAPTER\s+[IVXLCDM]+|(?:FIRST|SECOND|THIRD|FOURTH|FIFTH|SIXTH|SEVENTH|EIGHTH)\s+SCHEDULE))' |
| body_blocks = re.split(split_pattern, text) |
| |
| splitter = RecursiveCharacterTextSplitter( |
| chunk_size=800, |
| chunk_overlap=80, |
| length_function=len, |
| separators=["\n\n", "\n", " ", ""] |
| ) |
| |
| for block in body_blocks: |
| block = block.strip() |
| if not block: |
| continue |
| |
| lines = block.split('\n') |
| header = "" |
| if lines and re.search(r'^(Section\s+\d+(?:\.\d+)*|Rule\s+\d+(?:\.\d+)*|CHAPTER\s+[IVXLCDM]+|(?:FIRST|SECOND|THIRD|FOURTH|FIFTH|SIXTH|SEVENTH|EIGHTH)\s+SCHEDULE)', lines[0].strip(), re.IGNORECASE): |
| header = lines[0].strip() |
| |
| sub_chunks = splitter.split_text(block) |
| |
| for sc in sub_chunks: |
| sc = sc.strip() |
| if not sc: |
| continue |
| |
| if header and not sc.startswith(header.split('\n')[0]): |
| sc = header + "\n" + sc |
| |
| chunks.append(sc) |
| metadatas.append(extract_metadata(header, sc)) |
| |
| return chunks, metadatas |
|
|
|
|
| |
| |
| |
| def export_rust_artifacts( |
| chunks: List[str], |
| metadatas: List[Dict[str, str]], |
| bm25: BM25Okapi, |
| dense_embeddings: np.ndarray, |
| safetensors_path: Path, |
| index_json_path: Path |
| ): |
| print("\n📦 Exporting Language-Agnostic Rust Artifacts...") |
|
|
| |
| if not HAS_SAFETENSORS: |
| raise ImportError("safetensors library is missing. Install with: pip install safetensors") |
|
|
| dense_f32 = np.ascontiguousarray(dense_embeddings, dtype=np.float32) |
| tensors_dict = { |
| "dense_embeddings": dense_f32 |
| } |
| |
| safetensors_path.parent.mkdir(parents=True, exist_ok=True) |
| save_file(tensors_dict, str(safetensors_path)) |
| print(f" • Dense Safetensors: {safetensors_path} ({os.path.getsize(safetensors_path)/(1024*1024):.2f} MB, shape={dense_f32.shape})") |
|
|
| |
| |
| idf_dict = {k: float(v) for k, v in bm25.idf.items()} |
| doc_lens = [int(dl) for dl in bm25.doc_len] |
|
|
| json_payload = { |
| "schema_version": "2.0.0", |
| "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), |
| "embedding_model": "BAAI/bge-small-en-v1.5", |
| "embedding_dimensions": int(dense_embeddings.shape[1]), |
| "total_chunks": len(chunks), |
| "bm25_params": { |
| "k1": float(bm25.k1), |
| "b": float(bm25.b), |
| "avgdl": float(bm25.avgdl), |
| "doc_len": doc_lens, |
| "idf": idf_dict |
| }, |
| "chunks": chunks, |
| "metadatas": metadatas |
| } |
|
|
| with open(index_json_path, "w", encoding="utf-8") as f: |
| json.dump(json_payload, f, ensure_ascii=False, indent=2) |
| print(f" • Hybrid JSON Index: {index_json_path} ({os.path.getsize(index_json_path)/(1024*1024):.2f} MB)") |
|
|
| |
| loaded_tensors = load_file(str(safetensors_path)) |
| assert "dense_embeddings" in loaded_tensors, "Safetensors key validation failed!" |
| assert loaded_tensors["dense_embeddings"].shape == dense_f32.shape, "Tensor shape mismatch!" |
| print(" ✅ Rust Artifact Integrity Verification Passed.") |
|
|
|
|
| |
| |
| |
| def build_db(): |
| print("═══════════════════════════════════════════════════════════════════════") |
| print("🚀 SOTA HYBRID SEARCH ENGINE & TRIPLE-FORMAT ARTIFACT BUILDER") |
| print("═══════════════════════════════════════════════════════════════════════") |
| |
| if not LAW_TEXT_PATH.exists(): |
| print(f"❌ Error: Legal text file not found at {LAW_TEXT_PATH}") |
| return |
|
|
| print("✂️ Chunking statutory corpus with strict header binding...") |
| chunks, metadatas = parse_and_chunk(LAW_TEXT_PATH) |
| print(f"✅ Created {len(chunks)} legal chunks (800-char envelope with header preservation).") |
|
|
| print("\n🧠 Computing normalized dense embeddings with BAAI/bge-small-en-v1.5...") |
| embed_model = SentenceTransformer("BAAI/bge-small-en-v1.5") |
| embeddings = embed_model.encode(chunks, batch_size=32, show_progress_bar=True, normalize_embeddings=True) |
| embeddings = np.array(embeddings, dtype=np.float32) |
|
|
| print("\n📚 Compiling BM25 Okapi Lexical Inverted Index...") |
| tokenize = get_standard_tokenizer() |
| tokenized_corpus = [tokenize(c) for c in chunks] |
| bm25 = BM25Okapi(tokenized_corpus, k1=1.5, b=0.75) |
|
|
| print("\n📥 Verifying Cross-Encoder Local Artifacts (BAAI/bge-reranker-v2-m3)...") |
| snapshot_download("BAAI/bge-reranker-v2-m3") |
| print("✅ Cross-Encoder verified locally.") |
|
|
| |
| |
| |
| print(f"\n💾 Serializing Python Hybrid Index to: {HYBRID_INDEX_PATH}...") |
| out_dict = { |
| "chunks": chunks, |
| "metadatas": metadatas, |
| "bm25_index": bm25, |
| "dense_embeddings": embeddings, |
| "embedding_model_name": "BAAI/bge-small-en-v1.5", |
| "chunk_count": len(chunks), |
| "build_timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat() |
| } |
| |
| HYBRID_INDEX_PATH.parent.mkdir(parents=True, exist_ok=True) |
| with open(HYBRID_INDEX_PATH, "wb") as f: |
| pickle.dump(out_dict, f, protocol=pickle.HIGHEST_PROTOCOL) |
| |
| print(f"✅ Pickle artifact successfully saved ({os.path.getsize(HYBRID_INDEX_PATH) / (1024*1024):.2f} MB).") |
|
|
| |
| |
| |
| export_rust_artifacts( |
| chunks=chunks, |
| metadatas=metadatas, |
| bm25=bm25, |
| dense_embeddings=embeddings, |
| safetensors_path=SAFETENSORS_PATH, |
| index_json_path=INDEX_JSON_PATH |
| ) |
|
|
| |
| |
| |
| print("\n🔍 Executing Multi-Vector Retrieval Verification Suite...") |
| queries = [ |
| ("What is the penalty for data breach?", ["Schedule", "Section 33"], False), |
| ("verifiable consent for children", ["Section 9", "Rule 10"], False), |
| ("algorithmic software audit obligations", ["Rule 13"], False), |
| ("Consent Manager interoperability", ["Rule 4", "First Schedule"], False), |
| ("cross-border data transfer restrictions", ["Section 16", "Rule 15"], False), |
| ("Standards for processing personal data by State instrumentalities", ["Second Schedule"], True) |
| ] |
| |
| all_passed = True |
| for q, expected_keywords, is_state_query in queries: |
| q_tokens = tokenize(q) |
| bm25_scores = bm25.get_scores(q_tokens) |
| bm25_ranks = np.argsort(bm25_scores)[::-1] |
| |
| q_emb = embed_model.encode([f"Represent this sentence for searching relevant passages: {q}"], normalize_embeddings=True)[0] |
| dense_scores = np.dot(embeddings, q_emb) |
| dense_ranks = np.argsort(dense_scores)[::-1] |
| |
| rrf_scores = np.zeros(len(chunks)) |
| for rank, idx in enumerate(bm25_ranks[:50]): |
| if not is_state_query and metadatas[idx].get("applies_to") == "state": |
| continue |
| rrf_scores[idx] += 1.0 / (60 + rank + 1) |
| |
| for rank, idx in enumerate(dense_ranks[:50]): |
| if not is_state_query and metadatas[idx].get("applies_to") == "state": |
| continue |
| rrf_scores[idx] += 1.0 / (60 + rank + 1) |
| |
| top_indices = np.argsort(rrf_scores)[::-1][:3] |
| |
| print(f"\nQuery: '{q}'") |
| for i, idx in enumerate(top_indices): |
| meta = metadatas[idx] |
| snippet = chunks[idx][:80].replace('\n', ' ') |
| print(f" Hit {i+1} [{meta['type']} {meta['number']} | Applies To: {meta['applies_to']}] -> {snippet}...") |
| |
| if not is_state_query and meta["applies_to"] == "state": |
| print(f" ❌ FATAL ERROR: Retrieved State-only chunk for private query!") |
| all_passed = False |
| |
| if all_passed: |
| print("\n🏆 ALL RETRIEVAL & ISOLATION CHECKS PASSED.") |
| else: |
| print("\n⚠️ RETRIEVAL CHECKS FAILED. Review metadata tagging and chunk boundaries.") |
|
|
|
|
| if __name__ == "__main__": |
| build_db() |
