File size: 2,039 Bytes
701cf7d | 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 | """
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()
|