Spaces:
Sleeping
Sleeping
File size: 2,636 Bytes
32e4125 |
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 |
#!/usr/bin/env python3
"""
Script to rebuild the vector database with embeddings locally.
Run this when you update the synthetic_policies documents.
"""
import logging
import sys
from pathlib import Path
# Add src to path so we can import modules
sys.path.insert(0, str(Path(__file__).parent / "src"))
def main():
"""Build embeddings for the corpus."""
logging.basicConfig(level=logging.INFO)
print("🔄 Building embeddings database...")
# Import after setting up path
from src.config import (
COLLECTION_NAME,
CORPUS_DIRECTORY,
DEFAULT_CHUNK_SIZE,
DEFAULT_OVERLAP,
EMBEDDING_DIMENSION,
EMBEDDING_MODEL_NAME,
RANDOM_SEED,
VECTOR_DB_PERSIST_PATH,
)
from src.ingestion.ingestion_pipeline import IngestionPipeline
from src.vector_store.vector_db import VectorDatabase
print(f"📁 Processing corpus: {CORPUS_DIRECTORY}")
print(f"🤖 Using model: {EMBEDDING_MODEL_NAME}")
print(f"📊 Target dimension: {EMBEDDING_DIMENSION}")
# Clear existing database
import shutil
if Path(VECTOR_DB_PERSIST_PATH).exists():
print(f"🗑️ Clearing existing database: {VECTOR_DB_PERSIST_PATH}")
shutil.rmtree(VECTOR_DB_PERSIST_PATH)
# Run ingestion pipeline
ingestion_pipeline = IngestionPipeline(
chunk_size=DEFAULT_CHUNK_SIZE,
overlap=DEFAULT_OVERLAP,
seed=RANDOM_SEED,
store_embeddings=True,
)
result = ingestion_pipeline.process_directory_with_embeddings(CORPUS_DIRECTORY)
chunks_processed = result["chunks_processed"]
embeddings_stored = result["embeddings_stored"]
if chunks_processed == 0:
print("❌ Ingestion failed or processed 0 chunks")
return 1
# Verify database
vector_db = VectorDatabase(VECTOR_DB_PERSIST_PATH, COLLECTION_NAME)
count = vector_db.get_count()
dimension = vector_db.get_embedding_dimension()
print(f"✅ Successfully processed {chunks_processed} chunks")
print(f"🔗 Embeddings stored: {embeddings_stored}")
print(f"📊 Database contains {count} embeddings")
print(f"🔢 Embedding dimension: {dimension}")
if dimension != EMBEDDING_DIMENSION:
print(f"⚠️ Warning: Expected dimension {EMBEDDING_DIMENSION}, got {dimension}")
return 1
print("🎉 Embeddings database ready for deployment!")
print("💡 Don't forget to commit the data/ directory to git")
# Clean up memory after build
import gc
gc.collect()
print("🧹 Memory cleanup completed")
return 0
if __name__ == "__main__":
sys.exit(main())
|