scripts/build_rag_db.py

"""
Build SQLite RAG knowledge base
"""
import sys
sys.path.append("src")
from ares.config import get_config
from ares.model.model import AresForCausalLM
from ares.tokenizer.tokenizer import AresTokenizer
from ares.memory.rag import RAGStore
import os

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="tiny")
    parser.add_argument("--tokenizer", type=str, default="data/tokenizer.json")
    parser.add_argument("--db", type=str, default="data/ares_knowledge.db")
    args = parser.parse_args()

    config = get_config(args.config)
    tokenizer = AresTokenizer(vocab_file=args.tokenizer if os.path.exists(args.tokenizer) else None, vocab_size=config.vocab_size)
    model = AresForCausalLM(config)
    device = "cuda" if __import__("torch").cuda.is_available() else "cpu"

    rag = RAGStore(model, tokenizer, db_path=args.db, device=device)
    print(f"[RAG] Building DB at {args.db}")
    # Ingest some knowledge
    sample_texts = [
        "Ares is named after Greek god of war, but in AI it's a general intelligence.",
        "Transformers architecture consists of attention, embeddings, RMSNorm, SwiGLU FFN, and unembedding.",
        "SQLite can handle terabytes with WAL mode and indexing.",
        "Chain-of-thought prompting improves reasoning by breaking tasks into steps.",
        "AdamW optimizer decouples weight decay from gradient-based update.",
        "RoPE encodes position via rotation, enabling long context extrapolation.",
        "GQA uses fewer KV heads than Q heads to save memory.",
        "KV Cache avoids recomputing past keys/values during autoregressive generation.",
        "Python code execution can be sandboxed with restricted builtins.",
        "RAG retrieves documents via cosine similarity of embeddings."
    ]
    rag.ingest(sample_texts, sources=["seed"]*len(sample_texts))
    print(rag.stats())
    # Optionally ingest HF wikitext
    print("[RAG] Stats:", rag.stats())

if __name__ == "__main__":
    main()