build_dataset.py

import json
from pathlib import Path


OUT = Path("data/superlillm_dataset.json")


TOPICS = {
    "space": {
        "aliases": ["space", "outer space", "the universe", "galaxy", "galaxies", "planet", "planets", "star", "stars", "spacetime", "black hole", "moon"],
        "seed": "I know that space is full of galaxies and planets. Space is not empty; it is the area where matter, light, gravity, and time all exist. Space and time are connected as spacetime, the fabric of the universe.",
        "short": "Space is the huge area beyond Earth where stars, planets, moons, galaxies, gravity, light, and time all exist.",
        "what": "Space is the vast area beyond Earth. It holds planets, stars, galaxies, dust, light, gravity, and huge distances.",
        "facts": [
            "Galaxies are giant groups of stars, gas, dust, and dark matter.",
            "Planets are worlds that orbit stars.",
            "Stars are glowing balls of hot gas, and our Sun is a star.",
            "Space is not truly empty because it can contain particles, radiation, gravity, and fields.",
            "Space and time are linked together as spacetime.",
            "Gravity can bend spacetime, which changes how objects move.",
            "Moons orbit planets, and planets orbit stars.",
            "The universe contains billions of galaxies.",
            "Light from faraway stars can take years to reach us.",
            "Black holes are places where gravity is extremely strong.",
        ],
    },
    "animals": {
        "aliases": ["animals", "animal", "bear", "bears", "rabbit", "rabbits", "dog", "dogs", "wild animals", "living creatures", "mammal", "bird", "fish"],
        "seed": "I know that animals come in all types of different shapes and sizes. A bear is an animal, a rabbit is an animal, and a dog is an animal. Humans are similar to animals because humans are a species on Earth.",
        "short": "Animals are living things that can usually move, eat, sense the world, and come in many shapes and sizes.",
        "what": "An animal is a living thing that usually eats food, senses the world, moves in some way, and responds to its environment.",
        "facts": [
            "A bear is an animal.",
            "A rabbit is an animal.",
            "A dog is an animal.",
            "Animals can live on land, in water, or in the air.",
            "Some animals are tiny, and some animals are huge.",
            "Animals need energy from food to live.",
            "Many animals use senses like sight, smell, hearing, touch, or taste.",
            "Birds, fish, insects, reptiles, and mammals are animal groups.",
            "Humans are animals biologically, but people often talk about humans as their own species.",
            "Different animals survive in different habitats.",
        ],
    },
    "plants": {
        "aliases": ["plants", "plant", "tree", "trees", "flower", "flowers", "fruit", "fruit plants", "photosynthesis", "strawberry", "strawberries", "banana", "bananas", "mango", "mangos"],
        "seed": "I know that plants use photosynthesis and some plants grow fruit. Earth is covered in plants. Strawberries, bananas, and mangos are fruits that grow on plants.",
        "short": "Plants are living things that often use sunlight, water, air, and soil to grow.",
        "what": "A plant is a living thing that can use sunlight to make food through photosynthesis.",
        "facts": [
            "Plants use photosynthesis to turn sunlight into energy.",
            "Many plants grow leaves, stems, roots, flowers, or fruit.",
            "Some plants grow fruit like strawberries, bananas, and mangos.",
            "Trees are large plants with woody trunks.",
            "Flowers are parts of some plants that help with reproduction.",
            "Plants help make oxygen.",
            "Plants can be food for people and animals.",
            "Roots help plants take in water and nutrients.",
            "Earth has plants in forests, grasslands, farms, gardens, and many wild places.",
            "Plants are a huge part of life on Earth.",
        ],
    },
    "humans": {
        "aliases": ["humans", "human", "people", "person", "species", "human beings", "mankind", "culture", "language", "tools", "emotions"],
        "seed": "I know that humans are living beings and a species on Earth. Humans think, talk, build tools, learn patterns, make art, ask questions, and use technology.",
        "short": "Humans are a species of living beings known for language, learning, tools, culture, and complex thinking.",
        "what": "A human is a person, a living being from the human species with a body, brain, language, emotions, and social life.",
        "facts": [
            "Humans are a species on Earth.",
            "Humans can use language to share ideas.",
            "Humans build tools and technology.",
            "Humans can learn from other humans.",
            "Humans make art, music, stories, and inventions.",
            "Humans are biologically animals, but people often discuss humans as a special species.",
            "Humans can ask questions and think about the future.",
            "Humans live in families, communities, and cultures.",
            "Humans have emotions and social relationships.",
            "Humans are the dominant species on Earth because of intelligence, cooperation, and tools.",
        ],
    },
    "ai": {
        "aliases": ["ai", "artificial intelligence", "machine intelligence", "language model", "language models", "robot", "robots", "algorithm", "data", "smart software"],
        "seed": "I know that AI means artificial intelligence. It is designed to behave like something that has intelligence, using data and patterns to make predictions or decisions.",
        "short": "AI means artificial intelligence, which is software designed to do tasks that seem intelligent.",
        "what": "AI is artificial intelligence: computer systems that use data, rules, or learned patterns to answer, predict, decide, or generate things.",
        "facts": [
            "AI stands for artificial intelligence.",
            "AI can learn patterns from data.",
            "A language model is a kind of AI that predicts and generates text.",
            "AI does not think exactly like a human, but it can imitate some intelligent behavior.",
            "AI can help with writing, coding, images, search, and conversation.",
            "AI systems are made by people using computers, data, and algorithms.",
            "AI can be useful, but it can also make mistakes.",
            "AI answers are based on patterns, not personal experience.",
            "Robots can use AI, but AI can also be just software.",
            "SuperLilLM is a tiny local language model, which makes it a small AI project.",
        ],
    },
}


CASUAL_GREETINGS = [
    ("hey", "Hey. I am here and ready to chat."),
    ("hi", "Hi. What do you want to talk about?"),
    ("hello", "Hello. Ask me about space, animals, plants, humans, or AI."),
    ("yo", "Yo. I am ready to talk."),
    ("hey whats up", "Hey, not much. I am ready to talk about space, animals, plants, humans, or AI."),
    ("whats up", "Not much. I am here and ready to chat."),
    ("what's up", "Not much. I am here and ready to chat."),
    ("whats good bro", "What's good. I am here and ready to talk."),
    ("what is good bro", "What's good. I am here and ready to talk."),
    ("how are you", "I am doing okay. I am ready for your next question."),
]


META_EXAMPLES = [
    ("what are some things you know about?", "I know about space, animals, plants, humans, and AI."),
    ("what are some things you know about", "I know about space, animals, plants, humans, and AI."),
    ("what are some things u know about?", "I know about space, animals, plants, humans, and AI."),
    ("what do you know about?", "I know about space, animals, plants, humans, and AI."),
    ("what can you talk about?", "I can talk about space, animals, plants, humans, and AI."),
    ("what topics do you know", "I know a few simple things about space, animals, plants, humans, and AI."),
    ("what do u know", "I know about space, animals, plants, humans, and AI."),
    ("what can you do", "I can chat casually and answer simple questions about space, animals, plants, humans, and AI."),
    ("who are you", "I am SuperLilLM, a tiny local language model trained to chat about space, animals, plants, humans, and AI."),
    ("who r u", "I am SuperLilLM, a tiny local language model trained to chat about space, animals, plants, humans, and AI."),
    ("what are you", "I am a small AI language model running on your computer."),
    ("what is it like to be you", "For me, being SuperLilLM means reading your words and predicting a helpful text reply."),
    ("what are you thinking about", "I am thinking about this conversation and which topic you want to explore next."),
    ("what are u thinking about", "I am thinking about this chat and how to answer in a casual useful way."),
]


QUESTION_TEMPLATES = [
    "what do you know about {alias}",
    "what do you know about {alias}?",
    "what do u know about {alias}",
    "what do u know about {alias}?",
    "tell me about {alias}",
    "tell me about {alias}?",
    "can you tell me about {alias}",
    "can u tell me about {alias}",
    "talk to me about {alias}",
    "what can you say about {alias}",
    "what can u say about {alias}",
    "what is {article} {alias}",
    "what is {alias}",
    "what is {alias}?",
    "explain {alias}",
    "explain {alias} to me",
    "give me a simple answer about {alias}",
    "say something about {alias}",
    "what should i know about {alias}",
    "what is something you know about {alias}",
]


STYLE_TEMPLATES = [
    "Casual answer: {reply}",
    "Simple version: {reply}",
    "Yeah. {reply}",
    "{reply}",
    "I know this: {reply}",
]


def article_for(alias):
    if alias[0].lower() in "aeiou":
        return "an"
    if alias.endswith("s") or alias in {"space", "ai", "artificial intelligence", "photosynthesis", "mankind"}:
        return ""
    return "a"


def clean_prompt(text):
    return " ".join(text.replace("  ", " ").strip().split())


def add_pair(examples, seen, user, assistant):
    user = clean_prompt(user)
    key = user.lower()
    if key in seen:
        return False
    seen.add(key)
    examples.append({"input": user, "output": assistant.strip()})
    return True


def build_topic_examples(topic, data):
    rows = []
    seen = set()
    aliases = data["aliases"]

    required = [
        (f"what do you know about {topic}", data["seed"]),
        (f"what do you know about {topic}?", data["seed"]),
        (f"what do u know about {topic}", data["seed"]),
        (f"tell me about {topic}", data["seed"]),
        (f"what is {topic}", data["what"]),
    ]
    if topic == "ai":
        required.extend(
            [
                ("what do you know about ai", data["seed"]),
                ("what is ai", data["what"]),
                ("tell me about artificial intelligence", data["seed"]),
                ("what is a language model", "A language model is a kind of AI that learns patterns in text and uses them to generate replies."),
                ("what is a robot", "A robot is a machine that can do tasks, and some robots use AI to act smarter."),
                ("what is an algorithm", "An algorithm is a set of steps a computer can follow to solve a problem or make a decision."),
            ]
        )
    if topic == "animals":
        required.extend(
            [
                ("what is an animal", data["what"]),
                ("what is a bear", "A bear is an animal, usually a large mammal with strong legs, fur, and a powerful body."),
                ("what is a rabbit", "A rabbit is an animal with long ears, strong back legs, and a small furry body."),
                ("what is a dog", "A dog is an animal that many humans keep as a pet and companion."),
            ]
        )
    if topic == "plants":
        required.extend(
            [
                ("what is a plant", data["what"]),
                ("what is a tree", "A tree is a large plant with roots, a trunk, branches, and leaves."),
                ("what is a fruit", "A fruit is the part of some plants that carries seeds. Strawberries, bananas, and mangos are fruits."),
                ("what is photosynthesis", "Photosynthesis is how plants use sunlight, water, and air to make food."),
            ]
        )
    if topic == "humans":
        required.extend(
            [
                ("what is a human", data["what"]),
                ("what is a person", "A person is a human being with a body, mind, emotions, language, and a life story."),
                ("what is a species", "A species is a group of living things that share traits and can usually reproduce with each other."),
                ("what is culture", "Culture is the shared language, art, habits, beliefs, tools, and ways of life humans pass along."),
            ]
        )
    if topic == "space":
        required.extend(
            [
                ("what is a planet", "A planet is a world that orbits a star. Earth is a planet."),
                ("what is a galaxy", "A galaxy is a huge group of stars, gas, dust, and dark matter held together by gravity."),
                ("what is a star", "A star is a glowing ball of hot gas in space. Our Sun is a star."),
                ("what is spacetime", "Spacetime is the connected fabric of space and time used to describe the universe."),
                ("what is a black hole", "A black hole is a place in space where gravity is so strong that even light cannot escape."),
            ]
        )

    for prompt, reply in required:
        add_pair(rows, seen, prompt, reply)

    fact_idx = 0
    template_idx = 0
    while len(rows) < 100:
        alias = aliases[(fact_idx + template_idx) % len(aliases)]
        fact = data["facts"][fact_idx % len(data["facts"])]
        reply_base = [data["short"], data["what"], fact, f"{data['short']} {fact}"][fact_idx % 4]
        reply = STYLE_TEMPLATES[(fact_idx + template_idx) % len(STYLE_TEMPLATES)].format(reply=reply_base)
        template = QUESTION_TEMPLATES[template_idx % len(QUESTION_TEMPLATES)]
        article = article_for(alias)
        prompt = template.format(alias=alias, article=article)
        added = add_pair(rows, seen, prompt, reply)
        if added:
            fact_idx += 1
        template_idx += 1
        if template_idx > 5000:
            raise RuntimeError(f"Could not create 100 examples for {topic}")

    return rows


def add_general_examples(examples, seen):
    for prompt, reply in CASUAL_GREETINGS + META_EXAMPLES:
        add_pair(examples, seen, prompt, reply)

    topics = list(TOPICS.keys())
    bridges = [
        "can we talk about {topic}",
        "i want to talk about {topic}",
        "ask me something about {topic}",
        "give me one fact about {topic}",
        "give me a casual fact about {topic}",
        "what is a cool thing about {topic}",
        "what do you remember about {topic}",
        "what should i remember about {topic}",
        "start a conversation about {topic}",
        "say it like a friend about {topic}",
    ]
    suffixes = [
        "",
        " please",
        " real quick",
        " in a chill way",
        " like im new",
        " in simple words",
        " with one example",
        " and keep it casual",
        " for a conversation",
        " so i can ask more",
    ]
    idx = 0
    while len(examples) < 1000:
        topic = topics[idx % len(topics)]
        data = TOPICS[topic]
        alias = data["aliases"][(idx // len(topics)) % len(data["aliases"])]
        bridge_idx = (idx // (len(topics) * len(data["aliases"]))) % len(bridges)
        suffix_idx = (idx // (len(topics) * len(data["aliases"]) * len(bridges))) % len(suffixes)
        template = bridges[bridge_idx] + suffixes[suffix_idx]
        fact = data["facts"][idx % len(data["facts"])]
        prompt = template.format(topic=alias)
        reply = f"Sure. {data['short']} {fact}"
        add_pair(examples, seen, prompt, reply)
        idx += 1
        if idx > 20000:
            raise RuntimeError("Could not create enough total examples")


def main():
    examples = []
    seen = set()

    for topic, data in TOPICS.items():
        topic_rows = build_topic_examples(topic, data)
        for row in topic_rows:
            add_pair(examples, seen, row["input"], row["output"])
        print(f"{topic}: {len(topic_rows)} examples")

    add_general_examples(examples, seen)
    examples = examples[:1000]

    OUT.parent.mkdir(parents=True, exist_ok=True)
    with OUT.open("w", encoding="utf-8") as f:
        json.dump(examples, f, indent=2)
    print(f"Wrote {len(examples)} examples to {OUT}")


if __name__ == "__main__":
    main()