File size: 1,860 Bytes
8eaa451 | 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 | """
Build ChromaDB offline - chay tren local, push len HF de khong bi timeout.
"""
import sys
import os
# Fix encoding cho Windows console
if sys.platform == "win32":
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
sys.stderr.reconfigure(encoding="utf-8", errors="replace")
import logging
logging.basicConfig(level=logging.INFO, stream=sys.stdout,
format="%(asctime)s - %(levelname)s - %(message)s")
# Add backend to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "backend")))
# Target paths
DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "backend", "data"))
CHROMA_DB_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "backend", "data", "chroma_db"))
# Override module-level constants BEFORE importing the function
import app.core.config as cfg
cfg.CHROMA_DB_PATH = CHROMA_DB_PATH
cfg.DATA_DIR = DATA_DIR
import app.services.chat_service as svc
from app.services.chat_service import _build_retrievers
svc.FORCE_REBUILD_DB = True
print("=" * 60)
print(f"DATA_DIR : {DATA_DIR}")
print(f"CHROMA_DB_PATH: {CHROMA_DB_PATH}")
print("=" * 60)
# Run the build
fast_ret, deep_ret, splits = _build_retrievers(DATA_DIR, CHROMA_DB_PATH)
# Verify
if os.path.exists(CHROMA_DB_PATH):
total_files = sum(len(files) for _, _, files in os.walk(CHROMA_DB_PATH))
total_size_mb = sum(os.path.getsize(os.path.join(r, f)) for r, _, files in os.walk(CHROMA_DB_PATH) for f in files) / (1024*1024)
print("=" * 60)
print(f"SUCCESS! ChromaDB created at: {CHROMA_DB_PATH}")
print(f" Files: {total_files}")
print(f" Size : {total_size_mb:.1f} MB")
print(f" Chunks indexed: {len(splits)}")
print("=" * 60)
else:
print("=" * 60)
print(f"ERROR: ChromaDB directory NOT found at {CHROMA_DB_PATH}")
print("=" * 60)
sys.exit(1)
|