""" Ingest Training Data Script ============================ Standalone script to convert all documents in ./train_data into embeddings and store them in the ChromaDB vector store. Usage: python ingest_train_data.py Requires: - AZURE_API_KEY environment variable set - data/config.json (or root config.json as fallback) with endpoint URLs - Documents (PDF/TXT) in ./train_data// """ import os import sys import json import logging import time from pathlib import Path os.environ.setdefault("ANONYMIZED_TELEMETRY", "False") import requests as http_requests import chromadb from chromadb.config import Settings from langchain_text_splitters import RecursiveCharacterTextSplitter from pypdf import PdfReader logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- PROJECT_ROOT = Path(__file__).parent DATA_DIR = Path("/data") if Path("/data").is_dir() else PROJECT_ROOT / "data" TRAIN_DATA_DIR = Path("/train_data") if Path("/train_data").is_dir() else PROJECT_ROOT / "train_data" CHROMA_PERSIST_DIR = str(DATA_DIR / "chroma_db") COLLECTION_NAME = "rag_documents" CHUNK_SIZE = 512 CHUNK_OVERLAP = 50 EMBEDDING_BATCH_SIZE = 16 _CONFIG_PATH = DATA_DIR / "config.json" if not _CONFIG_PATH.exists(): _CONFIG_PATH = PROJECT_ROOT / "config.json" logger.warning(f"No config.json in {DATA_DIR} — using root config.json as fallback.") with open(_CONFIG_PATH, encoding="utf-8") as _f: _config = json.load(_f) EMBEDDING_ENDPOINT_URL = _config["embedding"]["endpoint_url"] EMBEDDING_MODEL_NAME = _config["embedding"]["model"] AZURE_API_KEY = os.environ.get("AZURE_API_KEY") if not AZURE_API_KEY: logger.error("AZURE_API_KEY is not set. Export it before running this script.") sys.exit(1) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def extract_text_from_pdf(pdf_path: Path) -> str: reader = PdfReader(str(pdf_path)) pages_text = [] for page_num, page in enumerate(reader.pages, start=1): text = page.extract_text() if text and text.strip(): pages_text.append(f"[Page {page_num}]\n{text.strip()}") return "\n\n".join(pages_text) def chunk_text(text: str, source: str) -> list[dict]: splitter = RecursiveCharacterTextSplitter( chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, separators=["\n\n", "\n", ". ", " ", ""], ) chunks = splitter.split_text(text) return [{"text": chunk, "source": source, "chunk_index": i} for i, chunk in enumerate(chunks)] def generate_embeddings(texts: list[str]) -> list[list[float]]: headers = {"api-key": AZURE_API_KEY, "Content-Type": "application/json"} payload = {"input": texts, "model": EMBEDDING_MODEL_NAME} resp = http_requests.post(EMBEDDING_ENDPOINT_URL, headers=headers, json=payload, timeout=120) resp.raise_for_status() data = resp.json() return [item["embedding"] for item in data["data"]] def generate_embeddings_batched(texts: list[str]) -> list[list[float]]: all_embeddings = [] for i in range(0, len(texts), EMBEDDING_BATCH_SIZE): batch = texts[i:i + EMBEDDING_BATCH_SIZE] logger.info(f" Embedding batch {i // EMBEDDING_BATCH_SIZE + 1} ({len(batch)} chunks)...") embeddings = generate_embeddings(batch) all_embeddings.extend(embeddings) time.sleep(0.5) # Rate-limit courtesy return all_embeddings # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): if not TRAIN_DATA_DIR.exists(): logger.error(f"Training data directory not found: {TRAIN_DATA_DIR}") sys.exit(1) logger.info(f"Training data directory: {TRAIN_DATA_DIR}") logger.info(f"ChromaDB path: {CHROMA_PERSIST_DIR}") logger.info(f"Embedding model: {EMBEDDING_MODEL_NAME}") chroma_client = chromadb.PersistentClient( path=CHROMA_PERSIST_DIR, settings=Settings(anonymized_telemetry=False), ) collection = chroma_client.get_or_create_collection( name=COLLECTION_NAME, metadata={"hnsw:space": "cosine"}, ) logger.info(f"ChromaDB collection '{COLLECTION_NAME}' — existing documents: {collection.count()}") # Gather all PDF and TXT files recursively files = sorted( list(TRAIN_DATA_DIR.rglob("*.pdf")) + list(TRAIN_DATA_DIR.rglob("*.txt")) ) logger.info(f"Found {len(files)} files to process.") total_chunks_added = 0 for idx, file_path in enumerate(files, start=1): relative = file_path.relative_to(TRAIN_DATA_DIR) # Use category/filename as the source identifier source = str(relative) logger.info(f"[{idx}/{len(files)}] Processing: {source}") try: if file_path.suffix.lower() == ".pdf": text = extract_text_from_pdf(file_path) else: text = file_path.read_text(encoding="utf-8") except Exception as e: logger.warning(f" Skipped (read error): {e}") continue if not text.strip(): logger.warning(f" Skipped (no extractable text)") continue chunks = chunk_text(text, source=source) if not chunks: logger.warning(f" Skipped (no chunks produced)") continue logger.info(f" {len(chunks)} chunks extracted") texts = [c["text"] for c in chunks] try: embeddings = generate_embeddings_batched(texts) except http_requests.exceptions.HTTPError as e: logger.error(f" Embedding failed: {e}") continue except Exception as e: logger.error(f" Unexpected embedding error: {e}") continue existing_count = collection.count() ids = [f"doc_{existing_count + i}" for i in range(len(chunks))] metadatas = [{"source": c["source"], "chunk_index": c["chunk_index"]} for c in chunks] collection.add( ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas, ) total_chunks_added += len(chunks) logger.info(f" Added {len(chunks)} chunks (total in store: {collection.count()})") logger.info("=" * 60) logger.info(f"Ingestion complete. Chunks added this run: {total_chunks_added}") logger.info(f"Total documents in vector store: {collection.count()}") if __name__ == "__main__": main()