Blame for scripts/prepare_data.py

"""
ViuAI-1M Data Preparation Script.

Downloads and prepares a sample conversation dataset for training.
Creates JSONL files in the ViuAI chat format.

Usage:
    python scripts/prepare_data.py
"""

import json
import os
import random

# Output paths
DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data")
TRAIN_PATH = os.path.join(DATA_DIR, "train.jsonl")
EVAL_PATH = os.path.join(DATA_DIR, "eval.jsonl")
TOKENIZER_CORPUS_PATH = os.path.join(DATA_DIR, "tokenizer_corpus.txt")


def create_sample_conversations():
    """Create a sample conversation dataset for initial testing.

    In production, replace this with real conversation data from:
    - OpenAssistant Conversations (https://huggingface.co/datasets/OpenAssistant/oasst1)
    - ShareGPT (https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered)
    - Alpaca (https://github.com/tatsu-lab/stanford_alpaca)
    - Dolly (https://huggingface.co/datasets/databricks/databricks-dolly-15k)
    """
    conversations = [
        # Greetings
        {"messages": [
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hello! how can i help you today?"}
        ]},
        {"messages": [
            {"role": "user", "content": "hi there"},
            {"role": "assistant", "content": "hi! nice to meet you. what would you like to talk about?"}
        ]},
        {"messages": [
            {"role": "user", "content": "hey"},
            {"role": "assistant", "content": "hey! how are you doing today?"}
        ]},
        {"messages": [
            {"role": "user", "content": "good morning"},
            {"role": "assistant", "content": "good morning! hope you have a wonderful day. how can i assist you?"}
        ]},

        # General knowledge
        {"messages": [
            {"role": "user", "content": "what is machine learning?"},
            {"role": "assistant", "content": "machine learning is a type of artificial intelligence where computers learn from data to make predictions or decisions without being explicitly programmed for every task."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is python?"},
            {"role": "assistant", "content": "python is a popular programming language known for being easy to read and write. it is widely used for web development, data science, machine learning, and automation."}
        ]},
        {"messages": [
            {"role": "user", "content": "explain neural networks"},
            {"role": "assistant", "content": "neural networks are computing systems inspired by the human brain. they consist of layers of connected nodes that process information. each connection has a weight that adjusts during training to improve accuracy."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is deep learning?"},
            {"role": "assistant", "content": "deep learning is a subset of machine learning that uses neural networks with many layers. it excels at tasks like image recognition, natural language processing, and speech recognition."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is artificial intelligence?"},
            {"role": "assistant", "content": "artificial intelligence is the simulation of human intelligence by machines. it includes learning, reasoning, problem solving, perception, and language understanding."}
        ]},

        # Coding help
        {"messages": [
            {"role": "user", "content": "how do i print hello world in python?"},
            {"role": "assistant", "content": "you can print hello world in python by writing: print('hello world'). this will display the text on your screen."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is a variable in programming?"},
            {"role": "assistant", "content": "a variable is a named container that stores data in a program. you can think of it like a labeled box that holds a value. for example, in python: x = 5 creates a variable named x with the value 5."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is a function?"},
            {"role": "assistant", "content": "a function is a reusable block of code that performs a specific task. you define it once and can call it many times. this helps organize code and avoid repetition."}
        ]},
        {"messages": [
            {"role": "user", "content": "how do loops work?"},
            {"role": "assistant", "content": "loops repeat a block of code multiple times. a for loop runs a specific number of times, while a while loop runs until a condition becomes false. they help automate repetitive tasks."}
        ]},

        # Math
        {"messages": [
            {"role": "user", "content": "what is 2 plus 2?"},
            {"role": "assistant", "content": "2 plus 2 equals 4."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is the square root of 144?"},
            {"role": "assistant", "content": "the square root of 144 is 12."}
        ]},
        {"messages": [
            {"role": "user", "content": "explain algebra"},
            {"role": "assistant", "content": "algebra is a branch of mathematics that uses symbols and letters to represent numbers and quantities in formulas and equations. it helps solve problems by finding unknown values."}
        ]},

        # Science
        {"messages": [
            {"role": "user", "content": "what is gravity?"},
            {"role": "assistant", "content": "gravity is a force that pulls objects toward each other. on earth, it pulls everything toward the center of the planet. it is what keeps us on the ground and the moon orbiting earth."}
        ]},
        {"messages": [
            {"role": "user", "content": "how does the sun produce energy?"},
            {"role": "assistant", "content": "the sun produces energy through nuclear fusion. in its core, hydrogen atoms combine to form helium, releasing enormous amounts of energy as light and heat."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is dna?"},
            {"role": "assistant", "content": "dna stands for deoxyribonucleic acid. it is the molecule that carries genetic instructions for the development, functioning, growth, and reproduction of all living organisms."}
        ]},

        # Conversational
        {"messages": [
            {"role": "user", "content": "tell me a joke"},
            {"role": "assistant", "content": "why did the programmer quit his job? because he did not get arrays! (a raise)"}
        ]},
        {"messages": [
            {"role": "user", "content": "what is your name?"},
            {"role": "assistant", "content": "my name is viuai. i am a small language model built for research purposes. nice to meet you!"}
        ]},
        {"messages": [
            {"role": "user", "content": "who created you?"},
            {"role": "assistant", "content": "i was created by the viuai team as a research project to explore language model architectures. i am a small model with about 1 million parameters."}
        ]},
        {"messages": [
            {"role": "user", "content": "can you help me?"},
            {"role": "assistant", "content": "of course! i am here to help. please tell me what you need assistance with and i will do my best."}
        ]},
        {"messages": [
            {"role": "user", "content": "thank you"},
            {"role": "assistant", "content": "you are welcome! i am glad i could help. feel free to ask me anything else."}
        ]},
        {"messages": [
            {"role": "user", "content": "goodbye"},
            {"role": "assistant", "content": "goodbye! it was nice talking to you. have a great day!"}
        ]},

        # Multi-turn conversations
        {"messages": [
            {"role": "user", "content": "i want to learn programming"},
            {"role": "assistant", "content": "that is great! programming is a wonderful skill. what language would you like to start with? i recommend python for beginners."},
            {"role": "user", "content": "yes python sounds good"},
            {"role": "assistant", "content": "perfect! let us start with the basics. first, you need to install python from python.org. then you can write your first program: print('hello world'). this displays text on the screen."}
        ]},
        {"messages": [
            {"role": "user", "content": "explain transformers in ai"},
            {"role": "assistant", "content": "transformers are a type of neural network architecture introduced in 2017. they use a mechanism called self-attention to process all parts of the input simultaneously, making them very efficient."},
            {"role": "user", "content": "what is self-attention?"},
            {"role": "assistant", "content": "self-attention allows each part of the input to look at all other parts to understand context. for example, in the sentence 'the cat sat on the mat', self-attention helps the model understand that 'cat' and 'sat' are related."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is the weather like?"},
            {"role": "assistant", "content": "i am sorry, i do not have access to real-time weather data. i am a text-based language model and cannot check current conditions. you can check a weather website or app for accurate information."},
            {"role": "user", "content": "that is okay, what can you do?"},
            {"role": "assistant", "content": "i can help with answering questions, explaining concepts, having conversations, helping with coding, math, science, and general knowledge topics. i am a small research model, so my abilities are limited compared to larger models."}
        ]},

        # Technology
        {"messages": [
            {"role": "user", "content": "what is the internet?"},
            {"role": "assistant", "content": "the internet is a global network of connected computers that allows people to share information, communicate, and access services. it uses protocols like tcp/ip to transmit data across the world."}
        ]},
        {"messages": [
            {"role": "user", "content": "how do computers work?"},
            {"role": "assistant", "content": "computers work by processing instructions using a central processing unit (cpu). they store data in memory (ram) and on storage drives. programs tell the computer what to do step by step using binary code (ones and zeros)."}
        ]},
        {"messages": [
            {"role": "user", "content": "what is cloud computing?"},
            {"role": "assistant", "content": "cloud computing is delivering computing services like servers, storage, databases, and software over the internet. instead of owning physical hardware, you rent these resources from providers like aws, google cloud, or azure."}
        ]},

        # Philosophy and thinking
        {"messages": [
            {"role": "user", "content": "what is the meaning of life?"},
            {"role": "assistant", "content": "the meaning of life is a deep philosophical question that has been debated for centuries. different people find meaning through relationships, work, creativity, helping others, or spiritual beliefs. what matters most is finding what gives your life purpose and fulfillment."}
        ]},
        {"messages": [
            {"role": "user", "content": "can machines think?"},
            {"role": "assistant", "content": "this is a famous question posed by alan turing. current ai systems like me can process language and generate responses, but we do not truly think or feel. we pattern match based on training data. the question of machine consciousness remains open."}
        ]},
    ]

    return conversations


def augment_conversations(conversations, multiplier=5):
    """Create variations of conversations to increase dataset size.

    Simple augmentation: slight text variations and reordering.
    """
    augmented = list(conversations)

    variations = {
        "hello": ["hi", "hey", "greetings", "howdy"],
        "thank you": ["thanks", "thank you so much", "thanks a lot", "appreciated"],
        "goodbye": ["bye", "see you", "take care", "farewell"],
        "what is": ["explain", "describe", "tell me about", "define"],
        "how do": ["how can", "what is the way to", "how would"],
    }

    for _ in range(multiplier):
        for conv in conversations:
            new_messages = []
            for msg in conv["messages"]:
                content = msg["content"]
                # Apply random variations
                for original, alts in variations.items():
                    if original in content and random.random() > 0.5:
                        content = content.replace(original, random.choice(alts), 1)
                new_messages.append({"role": msg["role"], "content": content})
            augmented.append({"messages": new_messages})

    return augmented


def main():
    """Prepare the training dataset."""
    print("=" * 60)
    print("  ViuAI-1M Data Preparation")
    print("=" * 60)

    os.makedirs(DATA_DIR, exist_ok=True)

    # Create sample conversations
    print("\n[1/4] Creating sample conversations...")
    conversations = create_sample_conversations()
    print(f"  Base conversations: {len(conversations)}")

    # Augment data
    print("\n[2/4] Augmenting dataset...")
    augmented = augment_conversations(conversations, multiplier=8)
    random.shuffle(augmented)
    print(f"  Total conversations: {len(augmented)}")

    # Split into train/eval
    print("\n[3/4] Splitting into train/eval...")
    split_idx = int(len(augmented) * 0.9)
    train_data = augmented[:split_idx]
    eval_data = augmented[split_idx:]

    # Save JSONL files
    with open(TRAIN_PATH, "w", encoding="utf-8") as f:
        for conv in train_data:
            f.write(json.dumps(conv, ensure_ascii=False) + "\n")

    with open(EVAL_PATH, "w", encoding="utf-8") as f:
        for conv in eval_data:
            f.write(json.dumps(conv, ensure_ascii=False) + "\n")

    print(f"  Train: {len(train_data)} conversations → {TRAIN_PATH}")
    print(f"  Eval:  {len(eval_data)} conversations → {EVAL_PATH}")

    # Create tokenizer training corpus
    print("\n[4/4] Creating tokenizer corpus...")
    corpus_texts = []
    for conv in augmented:
        for msg in conv["messages"]:
            corpus_texts.append(msg["content"])

    with open(TOKENIZER_CORPUS_PATH, "w", encoding="utf-8") as f:
        for text in corpus_texts:
            f.write(text + "\n")

    print(f"  Corpus texts: {len(corpus_texts)} → {TOKENIZER_CORPUS_PATH}")

    print(f"\n{'='*60}")
    print("  Data preparation complete!")
    print(f"{'='*60}")
    print(f"\nNext steps:")
    print(f"  1. For better results, add more conversation data to {TRAIN_PATH}")
    print(f"  2. Recommended datasets:")
    print(f"     - OpenAssistant: huggingface.co/datasets/OpenAssistant/oasst1")
    print(f"     - Dolly:        huggingface.co/datasets/databricks/databricks-dolly-15k")
    print(f"     - Alpaca:       github.com/tatsu-lab/stanford_alpaca")
    print(f"  3. Run training: python scripts/train.py")


if __name__ == "__main__":
    main()