build_rag.py

"""
clankerDiffusion — build the RAG knowledge base + RAG training data.

Run AFTER the tokenizer exists (data/tokenizer.json):
    python build_rag.py

  * writes a small default knowledge base into data/kb/ and saves its index
  * writes data/rag_corpus.txt : many formatted examples teaching the model
      - to call <tool name="retrieve">q</tool> and use the <result>, and
      - to read <context>...</context> injected mid-response and answer from it.
    This corpus is consumed by rag_finetune.py (continuation training).
"""
import os
import json
import random

from rag import KnowledgeBase, DEFAULT_INDEX, KB_DIR

HERE = os.path.dirname(os.path.abspath(__file__))
DATADIR = os.path.join(HERE, "data")
CORPUS = os.path.join(DATADIR, "rag_corpus.txt")

# A small factual knowledge base. (In practice, point build_rag.py at your own
# docs via kb.ingest_path(...); this seed makes `retrieve` work out of the box.)
FACTS = [
    ("geography", "France is a country in Western Europe. Its capital and largest city is Paris. The currency is the euro.",
     "What is the capital of France?", "The capital of France is Paris."),
    ("geography", "Germany's capital is Berlin. The official language is German and the currency is the euro.",
     "What is the capital of Germany?", "The capital of Germany is Berlin."),
    ("geography", "Japan is an island nation in East Asia. Its capital is Tokyo and its highest mountain is Mount Fuji.",
     "What is the capital of Japan?", "The capital of Japan is Tokyo."),
    ("science", "Water is a chemical compound with the formula H2O. It boils at 100 degrees Celsius at standard pressure and freezes at 0 degrees Celsius.",
     "At what temperature does water boil?", "Water boils at 100 degrees Celsius at standard pressure."),
    ("science", "The speed of light in a vacuum is approximately 299,792 kilometers per second.",
     "What is the speed of light?", "The speed of light in a vacuum is about 299,792 km/s."),
    ("science", "The chemical symbol for gold is Au, for silver is Ag, and for iron is Fe.",
     "What is the chemical symbol for gold?", "The chemical symbol for gold is Au."),
    ("history", "World War II lasted from 1939 to 1945. It involved most of the world's nations and ended with the defeat of the Axis powers.",
     "When did World War II end?", "World War II ended in 1945."),
    ("history", "The Declaration of Independence of the United States was adopted on July 4, 1776.",
     "When was the US Declaration of Independence adopted?", "It was adopted on July 4, 1776."),
    ("tech", "Python is a high-level, interpreted programming language. It uses indentation to define blocks and is widely used for data science and AI.",
     "What kind of language is Python?", "Python is a high-level, interpreted programming language."),
    ("tech", "The Transformer is a neural network architecture introduced in 2017 that relies on self-attention instead of recurrence.",
     "What is a Transformer in machine learning?", "A Transformer is a neural architecture from 2017 based on self-attention."),
    ("math", "The value of pi is approximately 3.14159. It is the ratio of a circle's circumference to its diameter.",
     "What is the value of pi?", "Pi is approximately 3.14159."),
    ("math", "Euler's identity is e^(i*pi) + 1 = 0, linking the numbers e, i, pi, 1, and 0.",
     "What is Euler's identity?", "Euler's identity is e^(i*pi) + 1 = 0."),
    ("space", "The Sun is the star at the center of the Solar System. Earth orbits it at an average distance of about 149.6 million kilometers.",
     "What is the Sun?", "The Sun is the star at the center of the Solar System."),
    ("space", "Mars is the fourth planet from the Sun and is often called the Red Planet because of its iron-oxide surface.",
     "Why is Mars called the Red Planet?", "Mars is called the Red Planet due to its iron-oxide (rusty) surface."),
    ("biology", "DNA stands for deoxyribonucleic acid. It carries the genetic instructions used in the growth and functioning of all known living organisms.",
     "What does DNA stand for?", "DNA stands for deoxyribonucleic acid."),
    ("biology", "Photosynthesis is the process by which plants convert light energy, water, and carbon dioxide into glucose and oxygen.",
     "What is photosynthesis?", "Photosynthesis is how plants turn light, water, and CO2 into glucose and oxygen."),
    ("economics", "Inflation is the rate at which the general level of prices for goods and services rises, eroding purchasing power.",
     "What is inflation?", "Inflation is the rise in the general price level, reducing purchasing power."),
    ("economics", "Gross Domestic Product (GDP) is the total monetary value of all finished goods and services produced within a country in a period.",
     "What is GDP?", "GDP is the total value of finished goods and services produced in a country."),
    ("language", "The word 'clanker' in this project is the name of the assistant. It is a from-scratch hybrid diffusion language model.",
     "What is clanker?", "clanker is the name of this assistant, a from-scratch hybrid diffusion language model."),
    ("project", "clankerDiffusion alternates between AR mode (normal next-token) and DIFF mode (masked diffusion) using a learned mode embedding.",
     "What are the two modes of clankerDiffusion?", "AR mode (autoregressive) and DIFF mode (masked diffusion)."),
]

SYSTEM = ("You are clanker, a helpful assistant that can use tools and read context. "
          "When you need knowledge, call <tool name=\"retrieve\">query</tool>. "
          "If <context>...</context> is provided, answer using it.")


def build_kb():
    os.makedirs(KB_DIR, exist_ok=True)
    # a readable KB file so list_dir/retrieve demo works
    kb_text = "\n\n".join(f"[{topic}]\n{doc}" for topic, doc, _, _ in FACTS)
    with open(os.path.join(KB_DIR, "knowledge.txt"), "w", encoding="utf-8") as f:
        f.write(kb_text)
    kb = KnowledgeBase()
    kb.ingest_path(KB_DIR)
    kb.save(DEFAULT_INDEX)
    return kb


def build_corpus(n_each=400, seed=0):
    random.seed(seed)
    lines = []
    facts = FACTS
    for _ in range(n_each):
        topic, doc, q, a = random.choice(facts)
        # format A: model must RETRIEVE then answer
        lines.append(
            f"<system>{SYSTEM}</system><user>{q}</user>"
            f"<assistant><tool name=\"retrieve\">{q}</tool>"
            f"<result>{doc}</result>{a}</assistant>")
        # format B: context already injected mid-response, model answers from it
        lines.append(
            f"<system>{SYSTEM}</system><user>{q}</user>"
            f"<assistant><context>{doc}</context>{a}</assistant>")
        # format C: a short chain-of-thought variant
        lines.append(
            f"<system>{SYSTEM}</system><user>{q}</user>"
            f"<assistant><think>The relevant knowledge is: {doc}</think>"
            f"{a}</assistant>")
    random.shuffle(lines)
    with open(CORPUS, "w", encoding="utf-8") as f:
        for ln in lines:
            f.write(ln + "\n")
    print(f"[rag] wrote {len(lines)} training examples -> {CORPUS}")


if __name__ == "__main__":
    build_kb()
    build_corpus()