first-aid-rag-assistant / src / ingest.py
ingest.py
Raw
import json
from pathlib import Path

import numpy as np
from sqlitesearch import TextSearchIndex, VectorSearchIndex
from tqdm.auto import tqdm

# Project root is two levels up from src/ingest.py (src/ → project root)
_ROOT = Path(__file__).parent.parent


def load_faq_data():
    documents = []

    with open(_ROOT / "data" / "firstaidqa_v1-first-half.json", "r") as file:
        documents = json.load(file)

    return documents


def build_index(documents):
    db_path = _ROOT / "sqlite" / "text" / "text-search.db"
    db_path.parent.mkdir(parents=True, exist_ok=True)
    index = TextSearchIndex(
        text_fields=["question", "answer"],
        db_path=str(db_path),
    )

    if index:
        return index
    else:
        index.fit(documents)

    return index


def build_vector_index(embedder, documents):
    db_path = _ROOT / "sqlite" / "vector" / "vector-search.db"
    db_path.parent.mkdir(parents=True, exist_ok=True)
    index = VectorSearchIndex(
        mode="hnsw",
        db_path=str(db_path),
    )

    if index:
        return index
    else:
        faq_texts = [f"Q: {item['question']} A: {item['answer']}" for item in documents]

        # Track indexing speed
        print("Generating Embeddings in Batches...")

        batch_size = 50
        db_vectors = []

        for i in tqdm(range(0, len(faq_texts), batch_size)):
            batch = faq_texts[i : i + batch_size]
            batch_vectors = embedder.encode_batch(batch)
            db_vectors.extend(batch_vectors)

        db_vectors = np.array(db_vectors)

        index.fit(db_vectors, documents)

    return index