Initial upload: ViuAI-1M architecture (1.18M params)

README.md

@@ -0,0 +1,141 @@
+---
+license: apache-2.0
+language:
+  - en
+tags:
+  - text-generation
+  - transformer
+  - from-scratch
+  - research
+  - viuai
+  - custom-architecture
+library_name: pytorch
+pipeline_tag: text-generation
+model-index:
+  - name: ViuAI-1M
+    results: []
+---
+
+# 🚀 ViuAI-1M
+
+**A lightweight ~1.18M parameter chat language model built entirely from scratch.**
+
+ViuAI-1M is a research-grade decoder-only transformer designed for learning and experimentation. Every component — from the BPE tokenizer to the attention mechanism — is custom-built with modern innovations.
+
+## ✨ Architecture Highlights
+
+| Component | Innovation | Benefit |
+|---|---|---|
+| **Attention** | Rotary Position Embeddings (RoPE) | Better length generalization |
+| **FFN** | SwiGLU activation | More expressive than ReLU/GELU |
+| **Normalization** | RMSNorm (Pre-Norm) | Faster & more stable training |
+| **Embeddings** | Weight-tied input/output | Fewer parameters |
+| **Tokenizer** | Custom BPE from scratch | Full vocabulary control |
+
+## 📊 Model Specifications
+
+| Spec | Value |
+|---|---|
+| Parameters | **1,180,800 (~1.18M)** |
+| Hidden Size | 128 |
+| Layers | 4 |
+| Attention Heads | 4 |
+| Head Dimension | 32 |
+| FFN (SwiGLU) | 256 |
+| Context Window | 512 tokens |
+| Vocabulary | 4,096 (BPE) |
+| Model Size | ~4.5 MB |
+
+## 🏗️ Architecture
+
+```
+Input Token IDs
+    → Token Embedding (4096 × 128)
+    → Transformer Block ×4
+        ├── RMSNorm → Multi-Head Attention (RoPE) → Residual
+        └── RMSNorm → SwiGLU FFN → Residual
+    → Final RMSNorm
+    → LM Head (weight-tied)
+    → Output Logits
+```
+
+## 💬 Chat Format
+
+```
+<bos><|system|>System prompt<|user|>User message<|assistant|>Response<eos>
+```
+
+## 🚀 Quick Start
+
+```python
+import torch
+from viuai.model.config import ViuAIConfig
+from viuai.model.viuai_model import ViuAIForCausalLM
+from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
+
+# Load config and model
+config = ViuAIConfig()
+model = ViuAIForCausalLM(config)
+print(f"Parameters: {model.count_parameters():,}")
+
+# Train tokenizer on your data
+tokenizer = ViuAITokenizer(vocab_size=4096)
+tokenizer.train(["your training texts here"])
+
+# Training
+optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+```
+
+## 📁 Repository Structure
+
+```
+├── viuai/
+│   ├── model/
+│   │   ├── config.py         # Model configuration
+│   │   ├── attention.py      # Multi-Head Attention + RoPE
+│   │   ├── layers.py         # RMSNorm + SwiGLU
+│   │   ├── transformer.py    # Transformer block
+│   │   └── viuai_model.py    # Full model + generation
+│   ├── tokenizer/
+│   │   └── bpe_tokenizer.py  # Custom BPE tokenizer
+│   ├── training/
+│   │   ├── dataset.py        # Dataset & data collation
+│   │   └── trainer.py        # Training loop & scheduler
+│   └── inference/
+│       └── chat.py           # Chat inference engine
+├── configs/
+│   └── viuai_1m.json         # Model hyperparameters
+├── scripts/
+│   ├── prepare_data.py       # Data preparation
+│   ├── train.py              # Training script
+│   └── chat.py               # Chat interface
+└── requirements.txt
+```
+
+## 📈 Training
+
+```bash
+# 1. Prepare data
+python scripts/prepare_data.py
+
+# 2. Train the model
+python scripts/train.py --epochs 20 --batch_size 32 --lr 3e-4
+
+# 3. Chat with the model
+python scripts/chat.py --checkpoint output/final --tokenizer output/tokenizer
+```
+
+## 🔗 Related Models
+
+- [ViuAI-500M](https://huggingface.co/ViuAI/ViuAT-500M) — 500M parameter model
+- [Sarus-500M](https://huggingface.co/ViuAI/Sarus-500M) — 500M text generation
+- [ViuMini-Dense-360M](https://huggingface.co/ViuAI/ViuMini-Dense-360M) — Dense 360M model
+- [Viu-1.5B-MoE](https://huggingface.co/ViuAI/Viu-1.5B-MoE) — 1.5B Mixture of Experts
+
+## 📄 License
+
+Apache 2.0
+
+---
+
+**Built with ❤️ by [ViuAI](https://huggingface.co/ViuAI)**

configs/viuai_1m.json

@@ -0,0 +1,31 @@
+{
+  "model_name": "ViuAI-1M",
+  "version": "0.1.0",
+  "architecture": {
+    "vocab_size": 4096,
+    "hidden_size": 128,
+    "num_layers": 4,
+    "num_heads": 4,
+    "head_dim": 32,
+    "intermediate_size": 256,
+    "max_seq_len": 512,
+    "dropout": 0.1,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-6,
+    "tie_word_embeddings": true
+  },
+  "special_tokens": {
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2
+  },
+  "training": {
+    "learning_rate": 3e-4,
+    "weight_decay": 0.01,
+    "warmup_steps": 100,
+    "max_grad_norm": 1.0,
+    "num_epochs": 10,
+    "batch_size": 32,
+    "eval_split": 0.1
+  }
+}

requirements.txt

@@ -0,0 +1,3 @@
+# ViuAI-1M Requirements
+torch>=2.0.0
+numpy>=1.24.0

scripts/chat.py

@@ -0,0 +1,57 @@
+"""
+ViuAI-1M Chat Script.
+
+Start an interactive chat session with a trained ViuAI model.
+
+Usage:
+    python scripts/chat.py --checkpoint output/final --tokenizer output/tokenizer
+"""
+
+import argparse
+import os
+import sys
+
+# Add project root to path
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, project_root)
+
+from viuai.inference.chat import load_model_for_chat
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Chat with ViuAI-1M")
+    parser.add_argument("--checkpoint", type=str, default="output/final",
+                        help="Path to model checkpoint directory")
+    parser.add_argument("--tokenizer", type=str, default="output/tokenizer",
+                        help="Path to tokenizer directory")
+    parser.add_argument("--device", type=str, default=None,
+                        help="Device (cuda/cpu)")
+    parser.add_argument("--temperature", type=float, default=0.8,
+                        help="Sampling temperature")
+    parser.add_argument("--max_tokens", type=int, default=128,
+                        help="Maximum tokens to generate")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    checkpoint_path = os.path.join(project_root, args.checkpoint)
+    tokenizer_path = os.path.join(project_root, args.tokenizer)
+
+    print(f"\n  Loading ViuAI-1M...")
+    print(f"  Checkpoint: {checkpoint_path}")
+    print(f"  Tokenizer:  {tokenizer_path}")
+
+    engine = load_model_for_chat(
+        checkpoint_dir=checkpoint_path,
+        tokenizer_dir=tokenizer_path,
+        device=args.device,
+    )
+
+    # Start interactive chat
+    engine.interactive()
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_data.py

@@ -0,0 +1,285 @@
+"""
+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()

scripts/train.py

@@ -0,0 +1,166 @@
+"""
+ViuAI-1M Training Script.
+
+End-to-end training pipeline:
+1. Load or prepare data
+2. Train BPE tokenizer on the corpus
+3. Create datasets and dataloaders
+4. Initialize model
+5. Train with the ViuAI trainer
+6. Save the final model and tokenizer
+
+Usage:
+    python scripts/train.py
+    python scripts/train.py --epochs 20 --batch_size 16 --lr 1e-4
+"""
+
+import argparse
+import os
+import sys
+import json
+
+# Add project root to path
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, project_root)
+
+import torch
+from torch.utils.data import DataLoader
+
+from viuai.model.config import ViuAIConfig
+from viuai.model.viuai_model import ViuAIForCausalLM
+from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
+from viuai.training.dataset import ViuAIDataset, ChatDataCollator
+from viuai.training.trainer import ViuAITrainer
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train ViuAI-1M")
+    parser.add_argument("--config", type=str, default="configs/viuai_1m.json",
+                        help="Path to model config JSON")
+    parser.add_argument("--train_data", type=str, default="data/train.jsonl",
+                        help="Path to training data JSONL")
+    parser.add_argument("--eval_data", type=str, default="data/eval.jsonl",
+                        help="Path to evaluation data JSONL")
+    parser.add_argument("--tokenizer_corpus", type=str, default="data/tokenizer_corpus.txt",
+                        help="Path to tokenizer training corpus")
+    parser.add_argument("--output_dir", type=str, default="output",
+                        help="Output directory for checkpoints and tokenizer")
+    parser.add_argument("--epochs", type=int, default=10,
+                        help="Number of training epochs")
+    parser.add_argument("--batch_size", type=int, default=32,
+                        help="Training batch size")
+    parser.add_argument("--lr", type=float, default=3e-4,
+                        help="Learning rate")
+    parser.add_argument("--warmup_steps", type=int, default=100,
+                        help="Number of warmup steps")
+    parser.add_argument("--device", type=str, default=None,
+                        help="Device (cuda/cpu/mps)")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"\n{'='*60}")
+    print(f"  ViuAI-1M Training Pipeline")
+    print(f"{'='*60}\n")
+
+    # ── Step 1: Load Config ──────────────────────────────────────
+    print("[1/6] Loading configuration...")
+    config_path = os.path.join(project_root, args.config)
+    with open(config_path, "r") as f:
+        config_data = json.load(f)
+
+    config = ViuAIConfig(**config_data["architecture"])
+    print(f"  Config: {config_path}")
+    print(f"  Estimated parameters: {config.num_parameters:,}")
+
+    # ── Step 2: Train Tokenizer ──────────────────────────────────
+    print(f"\n[2/6] Training BPE tokenizer...")
+    tokenizer = ViuAITokenizer(vocab_size=config.vocab_size)
+
+    corpus_path = os.path.join(project_root, args.tokenizer_corpus)
+    if os.path.exists(corpus_path):
+        with open(corpus_path, "r", encoding="utf-8") as f:
+            corpus_texts = [line.strip() for line in f if line.strip()]
+        tokenizer.train(corpus_texts, verbose=True)
+    else:
+        print(f"  WARNING: Corpus not found at {corpus_path}")
+        print(f"  Run 'python scripts/prepare_data.py' first!")
+        return
+
+    # Update config vocab size to match actual tokenizer
+    config.vocab_size = tokenizer.vocab_size
+
+    # Save tokenizer
+    tokenizer_dir = os.path.join(project_root, args.output_dir, "tokenizer")
+    tokenizer.save(tokenizer_dir)
+
+    # ── Step 3: Create Datasets ──────────────────────────────────
+    print(f"\n[3/6] Creating datasets...")
+    train_path = os.path.join(project_root, args.train_data)
+    eval_path = os.path.join(project_root, args.eval_data)
+
+    train_dataset = ViuAIDataset(train_path, tokenizer, max_seq_len=config.max_seq_len)
+    eval_dataset = ViuAIDataset(eval_path, tokenizer, max_seq_len=config.max_seq_len)
+
+    collator = ChatDataCollator(pad_token_id=config.pad_token_id)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        collate_fn=collator,
+        num_workers=0,
+    )
+    eval_loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collator,
+        num_workers=0,
+    )
+
+    print(f"  Train batches: {len(train_loader)}")
+    print(f"  Eval batches:  {len(eval_loader)}")
+
+    # ── Step 4: Initialize Model ─────────────────────────────────
+    print(f"\n[4/6] Initializing ViuAI-1M model...")
+    model = ViuAIForCausalLM(config)
+    print(model.get_parameter_summary())
+
+    # ── Step 5: Train ────────────────────────────────────────────
+    print(f"\n[5/6] Starting training...")
+    checkpoint_dir = os.path.join(project_root, args.output_dir, "checkpoints")
+
+    trainer = ViuAITrainer(
+        model=model,
+        train_dataloader=train_loader,
+        eval_dataloader=eval_loader,
+        learning_rate=args.lr,
+        weight_decay=0.01,
+        warmup_steps=args.warmup_steps,
+        num_epochs=args.epochs,
+        save_dir=checkpoint_dir,
+        log_interval=10,
+        eval_interval=50,
+        device=args.device,
+    )
+
+    metrics = trainer.train()
+
+    # ── Step 6: Save Final Model ─────────────────────────────────
+    print(f"\n[6/6] Saving final model...")
+    final_dir = os.path.join(project_root, args.output_dir, "final")
+    os.makedirs(final_dir, exist_ok=True)
+
+    torch.save(model.state_dict(), os.path.join(final_dir, "model.pt"))
+    config.save(os.path.join(final_dir, "config.json"))
+
+    print(f"  Model saved to: {final_dir}")
+    print(f"\n  To chat with your model:")
+    print(f"    python scripts/chat.py --checkpoint {final_dir} --tokenizer {tokenizer_dir}")
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -0,0 +1,20 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="viuai",
+    version="0.1.0",
+    description="ViuAI-1M: A lightweight chat language model built from scratch",
+    author="ViuAI Team",
+    packages=find_packages(),
+    python_requires=">=3.9",
+    install_requires=[
+        "torch>=2.0.0",
+        "numpy>=1.24.0",
+    ],
+    entry_points={
+        "console_scripts": [
+            "viuai-train=scripts.train:main",
+            "viuai-chat=scripts.chat:main",
+        ],
+    },
+)

viuai/__init__.py

@@ -0,0 +1,14 @@
+"""
+ViuAI-1M: A lightweight chat language model built from scratch.
+
+Architecture: Decoder-only Transformer with modern innovations
+- RMSNorm (faster than LayerNorm)
+- Rotary Position Embeddings (RoPE)
+- SwiGLU activation function
+- Weight-tied embeddings
+
+Parameters: ~1.18M
+"""
+
+__version__ = "0.1.0"
+__model_name__ = "ViuAI-1M"

viuai/inference/__init__.py

@@ -0,0 +1,3 @@
+from viuai.inference.chat import ViuAIChatEngine
+
+__all__ = ["ViuAIChatEngine"]

viuai/inference/chat.py

@@ -0,0 +1,194 @@
+"""
+ViuAI-1M Chat Inference Engine.
+
+Provides a simple interactive chat interface for the trained ViuAI model.
+Handles message formatting, tokenization, generation, and response decoding.
+"""
+
+import torch
+from typing import List, Dict, Optional
+
+from viuai.model.viuai_model import ViuAIForCausalLM
+from viuai.model.config import ViuAIConfig
+from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
+
+
+class ViuAIChatEngine:
+    """Chat engine for interacting with a trained ViuAI model.
+
+    Manages conversation history, formats messages using the chat template,
+    and generates responses using the model.
+    """
+
+    def __init__(
+        self,
+        model: ViuAIForCausalLM,
+        tokenizer: ViuAITokenizer,
+        device: Optional[str] = None,
+        system_prompt: str = "You are ViuAI, a helpful and friendly AI assistant.",
+    ):
+        """Initialize the chat engine.
+
+        Args:
+            model: Trained ViuAI model.
+            tokenizer: Trained ViuAI tokenizer.
+            device: Device to run inference on.
+            system_prompt: Default system prompt for conversations.
+        """
+        self.tokenizer = tokenizer
+        self.system_prompt = system_prompt
+
+        # Auto-detect device
+        if device:
+            self.device = torch.device(device)
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
+
+        self.model = model.to(self.device)
+        self.model.eval()
+
+        # Conversation history
+        self.history: List[Dict[str, str]] = []
+
+    def chat(
+        self,
+        user_message: str,
+        max_new_tokens: int = 128,
+        temperature: float = 0.8,
+        top_k: int = 40,
+        top_p: float = 0.9,
+    ) -> str:
+        """Send a message and get a response.
+
+        Args:
+            user_message: The user's input message.
+            max_new_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            top_k: Top-K sampling parameter.
+            top_p: Nucleus sampling parameter.
+
+        Returns:
+            The assistant's response string.
+        """
+        # Add user message to history
+        self.history.append({"role": "user", "content": user_message})
+
+        # Build full conversation with system prompt
+        messages = [{"role": "system", "content": self.system_prompt}] + self.history
+
+        # Format and tokenize
+        formatted = self.tokenizer.format_chat(messages)
+        # Add the assistant turn marker so model generates a response
+        formatted += "<|assistant|>"
+
+        input_ids = self.tokenizer.encode(formatted, add_special_tokens=True)
+        # Remove the trailing EOS since we want the model to generate
+        if input_ids and input_ids[-1] == self.tokenizer.token_to_id.get("<eos>", 2):
+            input_ids = input_ids[:-1]
+
+        input_tensor = torch.tensor([input_ids], dtype=torch.long, device=self.device)
+
+        # Generate response
+        output_ids = self.model.generate(
+            input_ids=input_tensor,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+
+        # Decode only the newly generated tokens
+        generated_ids = output_ids[0, len(input_ids):].tolist()
+        response = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+
+        # Clean up response (stop at any special token markers)
+        for stop_token in ["<|user|>", "<|assistant|>", "<|system|>"]:
+            if stop_token in response:
+                response = response[:response.index(stop_token)]
+
+        response = response.strip()
+
+        # Add assistant response to history
+        self.history.append({"role": "assistant", "content": response})
+
+        return response
+
+    def reset(self) -> None:
+        """Clear conversation history."""
+        self.history = []
+
+    def interactive(self) -> None:
+        """Start an interactive chat session in the terminal.
+
+        Type 'quit' or 'exit' to end the session.
+        Type 'reset' to clear conversation history.
+        """
+        print(f"\n{'='*60}")
+        print(f"  ViuAI-1M Interactive Chat")
+        print(f"{'='*60}")
+        print(f"  Model: ViuAI-1M ({self.model.count_parameters():,} params)")
+        print(f"  Device: {self.device}")
+        print(f"  Commands: 'quit', 'exit', 'reset'")
+        print(f"{'='*60}\n")
+
+        while True:
+            try:
+                user_input = input("You: ").strip()
+            except (EOFError, KeyboardInterrupt):
+                print("\nGoodbye!")
+                break
+
+            if not user_input:
+                continue
+            if user_input.lower() in ("quit", "exit"):
+                print("Goodbye!")
+                break
+            if user_input.lower() == "reset":
+                self.reset()
+                print("[Conversation reset]\n")
+                continue
+
+            response = self.chat(user_input)
+            print(f"ViuAI: {response}\n")
+
+
+def load_model_for_chat(
+    checkpoint_dir: str,
+    tokenizer_dir: str,
+    device: Optional[str] = None,
+) -> ViuAIChatEngine:
+    """Load a trained model and tokenizer for chat inference.
+
+    Args:
+        checkpoint_dir: Path to the model checkpoint directory.
+        tokenizer_dir: Path to the tokenizer directory.
+        device: Device to use for inference.
+
+    Returns:
+        Ready-to-use ViuAIChatEngine instance.
+    """
+    import os
+
+    # Load config
+    config = ViuAIConfig.load(os.path.join(checkpoint_dir, "config.json"))
+
+    # Load tokenizer
+    tokenizer = ViuAITokenizer.load(tokenizer_dir)
+
+    # Update config vocab size to match tokenizer
+    config.vocab_size = tokenizer.vocab_size
+
+    # Load model
+    model = ViuAIForCausalLM(config)
+    model.load_state_dict(
+        torch.load(
+            os.path.join(checkpoint_dir, "model.pt"),
+            map_location="cpu",
+        )
+    )
+
+    print(f"[ViuAI] Model loaded: {model.count_parameters():,} parameters")
+
+    return ViuAIChatEngine(model=model, tokenizer=tokenizer, device=device)

viuai/model/__init__.py

@@ -0,0 +1,4 @@
+from viuai.model.config import ViuAIConfig
+from viuai.model.viuai_model import ViuAIForCausalLM
+
+__all__ = ["ViuAIConfig", "ViuAIForCausalLM"]

viuai/model/attention.py

@@ -0,0 +1,154 @@
+"""
+ViuAI-1M Attention Mechanism.
+
+Implements Multi-Head Self-Attention with Rotary Position Embeddings (RoPE).
+RoPE encodes position information directly into the attention computation,
+eliminating the need for separate positional embeddings.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+
+
+def precompute_rope_frequencies(
+    head_dim: int,
+    max_seq_len: int,
+    theta: float = 10000.0,
+    device: Optional[torch.device] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Precompute the cosine and sine frequencies for RoPE.
+
+    Args:
+        head_dim: Dimensionality of each attention head (must be even).
+        max_seq_len: Maximum sequence length to precompute for.
+        theta: Base frequency for the rotary embeddings.
+        device: Device to place the tensors on.
+
+    Returns:
+        Tuple of (cos, sin) tensors, each of shape (max_seq_len, head_dim).
+    """
+    assert head_dim % 2 == 0, "head_dim must be even for RoPE"
+
+    # Compute frequency bands: theta_i = 1 / (theta^(2i/d)) for i = 0..d/2-1
+    freq_indices = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+    frequencies = 1.0 / (theta ** (freq_indices / head_dim))
+
+    # Compute position * frequency for all positions
+    positions = torch.arange(max_seq_len, dtype=torch.float32, device=device)
+    angles = torch.outer(positions, frequencies)  # (max_seq_len, head_dim/2)
+
+    # Duplicate to match full head_dim: [cos(0), cos(0), cos(1), cos(1), ...]
+    cos_cached = torch.cos(angles).repeat(1, 2)  # (max_seq_len, head_dim)
+    sin_cached = torch.sin(angles).repeat(1, 2)  # (max_seq_len, head_dim)
+
+    return cos_cached, sin_cached
+
+
+def apply_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """Apply Rotary Position Embeddings to input tensor.
+
+    Uses the rotation formula:
+        x_rotated = x * cos + rotate_half(x) * sin
+
+    Args:
+        x: Input tensor of shape (batch, num_heads, seq_len, head_dim).
+        cos: Cosine frequencies of shape (seq_len, head_dim).
+        sin: Sine frequencies of shape (seq_len, head_dim).
+
+    Returns:
+        Rotated tensor of the same shape as input.
+    """
+    seq_len = x.shape[2]
+    cos = cos[:seq_len].unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, head_dim)
+    sin = sin[:seq_len].unsqueeze(0).unsqueeze(0)
+
+    # Rotate half: [-x1, x0, -x3, x2, ...] (swap pairs and negate alternating)
+    x_rotated = torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1)
+    x_rotated = x_rotated.reshape(x.shape)
+
+    return (x * cos) + (x_rotated * sin)
+
+
+class ViuAIAttention(nn.Module):
+    """Multi-Head Self-Attention with Rotary Position Embeddings.
+
+    This implements the core attention mechanism of ViuAI-1M with:
+    - Separate Q, K, V projections (no bias for efficiency)
+    - RoPE applied to Q and K before attention computation
+    - Causal masking for autoregressive generation
+    - Dropout for regularization during training
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.head_dim = config.head_dim
+        self.hidden_size = config.hidden_size
+        self.dropout = config.dropout
+
+        # Linear projections without bias (standard in modern LLMs)
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        self.attn_dropout = nn.Dropout(self.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cos: torch.Tensor,
+        rope_sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for multi-head attention.
+
+        Args:
+            x: Input tensor of shape (batch_size, seq_len, hidden_size).
+            rope_cos: RoPE cosine frequencies.
+            rope_sin: RoPE sine frequencies.
+            attention_mask: Optional causal mask.
+
+        Returns:
+            Output tensor of shape (batch_size, seq_len, hidden_size).
+        """
+        batch_size, seq_len, _ = x.shape
+
+        # Project to Q, K, V and reshape for multi-head attention
+        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        # Shape: (batch_size, num_heads, seq_len, head_dim)
+
+        # Apply Rotary Position Embeddings to Q and K
+        q = apply_rope(q, rope_cos, rope_sin)
+        k = apply_rope(k, rope_cos, rope_sin)
+
+        # Scaled dot-product attention
+        scale = math.sqrt(self.head_dim)
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / scale
+        # Shape: (batch_size, num_heads, seq_len, seq_len)
+
+        # Apply causal mask (prevent attending to future tokens)
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        # Weighted sum of values
+        attn_output = torch.matmul(attn_weights, v)
+        # Shape: (batch_size, num_heads, seq_len, head_dim)
+
+        # Reshape back to (batch_size, seq_len, hidden_size)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_size)
+
+        # Final linear projection
+        return self.o_proj(attn_output)

viuai/model/config.py

@@ -0,0 +1,100 @@
+"""
+ViuAI-1M Model Configuration.
+
+Defines all hyperparameters for the ViuAI transformer architecture.
+"""
+
+import json
+from dataclasses import dataclass, field, asdict
+from typing import Optional
+
+
+@dataclass
+class ViuAIConfig:
+    """Configuration class for the ViuAI-1M model.
+
+    This config defines a ~1.18M parameter decoder-only transformer
+    with modern architectural innovations (RoPE, SwiGLU, RMSNorm).
+
+    Attributes:
+        vocab_size: Size of the tokenizer vocabulary.
+        hidden_size: Dimensionality of the hidden representations.
+        num_layers: Number of transformer decoder blocks.
+        num_heads: Number of attention heads.
+        head_dim: Dimensionality of each attention head.
+        intermediate_size: Inner dimensionality of the SwiGLU FFN.
+        max_seq_len: Maximum sequence length (context window).
+        dropout: Dropout probability for regularization.
+        rope_theta: Base frequency for Rotary Position Embeddings.
+        rms_norm_eps: Epsilon for RMSNorm numerical stability.
+        tie_word_embeddings: Whether to tie input/output embeddings.
+        bos_token_id: Beginning-of-sequence token ID.
+        eos_token_id: End-of-sequence token ID.
+        pad_token_id: Padding token ID.
+    """
+
+    # Architecture
+    vocab_size: int = 4096
+    hidden_size: int = 128
+    num_layers: int = 4
+    num_heads: int = 4
+    head_dim: int = 32
+    intermediate_size: int = 256
+    max_seq_len: int = 512
+    dropout: float = 0.1
+
+    # Positional encoding
+    rope_theta: float = 10000.0
+
+    # Normalization
+    rms_norm_eps: float = 1e-6
+
+    # Weight tying
+    tie_word_embeddings: bool = True
+
+    # Special tokens
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    pad_token_id: int = 0
+
+    def save(self, path: str) -> None:
+        """Save configuration to a JSON file."""
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(asdict(self), f, indent=2)
+
+    @classmethod
+    def load(cls, path: str) -> "ViuAIConfig":
+        """Load configuration from a JSON file."""
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls(**data)
+
+    def __post_init__(self):
+        assert self.hidden_size == self.num_heads * self.head_dim, (
+            f"hidden_size ({self.hidden_size}) must equal "
+            f"num_heads ({self.num_heads}) * head_dim ({self.head_dim})"
+        )
+
+    @property
+    def num_parameters(self) -> int:
+        """Estimate total parameter count."""
+        # Token embeddings
+        emb = self.vocab_size * self.hidden_size
+
+        # Per transformer layer
+        # Attention: Q, K, V projections + output projection (no bias)
+        attn = 4 * self.hidden_size * self.hidden_size
+        # SwiGLU FFN: gate + up + down projections (no bias)
+        ffn = 3 * self.hidden_size * self.intermediate_size
+        # RMSNorm: 2 per layer (attention + ffn), each has hidden_size params
+        norm = 2 * self.hidden_size
+        per_layer = attn + ffn + norm
+
+        # Final RMSNorm
+        final_norm = self.hidden_size
+
+        # LM head is tied with embeddings, so 0 extra
+        lm_head = 0 if self.tie_word_embeddings else self.vocab_size * self.hidden_size
+
+        total = emb + (self.num_layers * per_layer) + final_norm + lm_head
+        return total

viuai/model/layers.py

@@ -0,0 +1,68 @@
+"""
+ViuAI-1M Custom Layers.
+
+Implements the building blocks used in the ViuAI transformer:
+- RMSNorm: Root Mean Square Layer Normalization (faster than LayerNorm)
+- SwiGLU: Gated Linear Unit with Swish activation (better than ReLU/GELU)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization.
+
+    RMSNorm is computationally simpler than standard LayerNorm because it
+    does not compute the mean (no re-centering). It only normalizes by the
+    root mean square of the activations.
+
+    Formula: output = x / sqrt(mean(x^2) + eps) * weight
+
+    Reference: https://arxiv.org/abs/1910.07467
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Compute RMS
+        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+        # Normalize and scale
+        return (x / rms) * self.weight
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU Feed-Forward Network.
+
+    SwiGLU combines the Swish (SiLU) activation with a Gated Linear Unit.
+    It uses three linear projections instead of two:
+        1. Gate projection: computes the gating signal
+        2. Up projection: computes the value to be gated
+        3. Down projection: projects back to hidden_size
+
+    Formula: output = down_proj(silu(gate_proj(x)) * up_proj(x))
+
+    This is more expressive than standard FFN with ReLU/GELU and is used
+    in models like LLaMA, PaLM, and Gemini.
+
+    Reference: https://arxiv.org/abs/2002.05202
+    """
+
+    def __init__(self, hidden_size: int, intermediate_size: int, dropout: float = 0.1):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # SwiGLU: SiLU(gate) * up, then project down
+        gate = F.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        hidden = gate * up
+        hidden = self.dropout(hidden)
+        return self.down_proj(hidden)

viuai/model/transformer.py

@@ -0,0 +1,77 @@
+"""
+ViuAI-1M Transformer Block.
+
+Each transformer block consists of:
+1. RMSNorm → Multi-Head Self-Attention → Residual Connection
+2. RMSNorm → SwiGLU FFN → Residual Connection
+
+This follows the Pre-Norm architecture (norm before sublayer),
+which is more stable for training than Post-Norm.
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional
+
+from viuai.model.attention import ViuAIAttention
+from viuai.model.layers import RMSNorm, SwiGLU
+
+
+class ViuAITransformerBlock(nn.Module):
+    """A single transformer decoder block.
+
+    Implements the Pre-Norm transformer architecture with residual connections.
+    Each block processes the input through attention and feed-forward sublayers.
+    """
+
+    def __init__(self, config, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+
+        # Pre-norm for attention
+        self.attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Multi-head self-attention with RoPE
+        self.attention = ViuAIAttention(config)
+
+        # Pre-norm for feed-forward
+        self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # SwiGLU feed-forward network
+        self.feed_forward = SwiGLU(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            dropout=config.dropout,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cos: torch.Tensor,
+        rope_sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass through one transformer block.
+
+        Args:
+            x: Input tensor of shape (batch_size, seq_len, hidden_size).
+            rope_cos: RoPE cosine frequencies.
+            rope_sin: RoPE sine frequencies.
+            attention_mask: Optional causal attention mask.
+
+        Returns:
+            Output tensor of shape (batch_size, seq_len, hidden_size).
+        """
+        # Sublayer 1: Attention with residual
+        residual = x
+        x = self.attn_norm(x)
+        x = self.attention(x, rope_cos, rope_sin, attention_mask)
+        x = residual + x
+
+        # Sublayer 2: Feed-forward with residual
+        residual = x
+        x = self.ffn_norm(x)
+        x = self.feed_forward(x)
+        x = residual + x
+
+        return x

viuai/model/viuai_model.py

@@ -0,0 +1,291 @@
+"""
+ViuAI-1M: Full Causal Language Model.
+
+This is the main model class that assembles all components into a
+complete decoder-only transformer for autoregressive text generation.
+
+Architecture Overview:
+    Input IDs → Token Embedding → [Transformer Block × N] → RMSNorm → LM Head → Logits
+
+Key Features:
+    - ~1.18M parameters (lightweight for research)
+    - RoPE (Rotary Position Embeddings) — no separate position embedding layer
+    - SwiGLU activation in FFN
+    - RMSNorm instead of LayerNorm
+    - Weight-tied input/output embeddings
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Dict
+
+from viuai.model.config import ViuAIConfig
+from viuai.model.transformer import ViuAITransformerBlock
+from viuai.model.layers import RMSNorm
+from viuai.model.attention import precompute_rope_frequencies
+
+
+class ViuAIModel(nn.Module):
+    """The core ViuAI transformer model (without LM head).
+
+    This produces hidden states from input token IDs.
+    """
+
+    def __init__(self, config: ViuAIConfig):
+        super().__init__()
+        self.config = config
+
+        # Token embeddings (no position embeddings — RoPE handles position)
+        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+
+        # Embedding dropout
+        self.embedding_dropout = nn.Dropout(config.dropout)
+
+        # Stack of transformer blocks
+        self.layers = nn.ModuleList([
+            ViuAITransformerBlock(config, layer_idx=i)
+            for i in range(config.num_layers)
+        ])
+
+        # Final normalization
+        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Precompute RoPE frequencies (registered as buffer — not a parameter)
+        rope_cos, rope_sin = precompute_rope_frequencies(
+            head_dim=config.head_dim,
+            max_seq_len=config.max_seq_len,
+            theta=config.rope_theta,
+        )
+        self.register_buffer("rope_cos", rope_cos)
+        self.register_buffer("rope_sin", rope_sin)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass through the transformer.
+
+        Args:
+            input_ids: Token IDs of shape (batch_size, seq_len).
+            attention_mask: Optional attention mask.
+
+        Returns:
+            Hidden states of shape (batch_size, seq_len, hidden_size).
+        """
+        batch_size, seq_len = input_ids.shape
+
+        # Token embeddings
+        hidden_states = self.token_embedding(input_ids)
+        hidden_states = self.embedding_dropout(hidden_states)
+
+        # Create causal mask if not provided
+        if attention_mask is None:
+            attention_mask = self._create_causal_mask(seq_len, hidden_states.device)
+
+        # Pass through all transformer blocks
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                self.rope_cos,
+                self.rope_sin,
+                attention_mask,
+            )
+
+        # Final normalization
+        hidden_states = self.final_norm(hidden_states)
+
+        return hidden_states
+
+    @staticmethod
+    def _create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+        """Create a causal (lower-triangular) attention mask.
+
+        Future positions are masked with -inf so they get zero attention weight
+        after softmax.
+        """
+        mask = torch.triu(
+            torch.full((seq_len, seq_len), float("-inf"), device=device),
+            diagonal=1,
+        )
+        return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)
+
+
+class ViuAIForCausalLM(nn.Module):
+    """ViuAI-1M model with a language modeling head for text generation.
+
+    This wraps the core transformer with a linear projection (LM head)
+    that maps hidden states to vocabulary logits for next-token prediction.
+    """
+
+    def __init__(self, config: ViuAIConfig):
+        super().__init__()
+        self.config = config
+        self.model = ViuAIModel(config)
+
+        # LM Head: projects hidden states to vocabulary logits
+        if config.tie_word_embeddings:
+            # Share weights between token embedding and LM head
+            self.lm_head = None  # Will use embedding weight directly
+        else:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        """Initialize weights using a scaled normal distribution.
+
+        This follows the common practice in transformer models where
+        weights are initialized from N(0, 0.02) and biases to zero.
+        """
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """Forward pass with optional loss computation.
+
+        Args:
+            input_ids: Token IDs of shape (batch_size, seq_len).
+            labels: Target token IDs for loss computation. Same shape as input_ids.
+                    Positions with value -100 are ignored in loss.
+            attention_mask: Optional attention mask.
+
+        Returns:
+            Dictionary with 'logits' and optionally 'loss'.
+        """
+        # Get hidden states from the transformer
+        hidden_states = self.model(input_ids, attention_mask)
+
+        # Project to vocabulary logits
+        if self.config.tie_word_embeddings:
+            logits = F.linear(hidden_states, self.model.token_embedding.weight)
+        else:
+            logits = self.lm_head(hidden_states)
+
+        output = {"logits": logits}
+
+        # Compute cross-entropy loss if labels are provided
+        if labels is not None:
+            # Shift: predict next token (logits[:-1] predicts labels[1:])
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            output["loss"] = loss
+
+        return output
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 128,
+        temperature: float = 0.8,
+        top_k: int = 40,
+        top_p: float = 0.9,
+        repetition_penalty: float = 1.1,
+        eos_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Generate text autoregressively using sampling.
+
+        Args:
+            input_ids: Initial token IDs of shape (batch_size, seq_len).
+            max_new_tokens: Maximum number of new tokens to generate.
+            temperature: Sampling temperature (higher = more random).
+            top_k: Keep only top-k tokens for sampling.
+            top_p: Nucleus sampling threshold.
+            repetition_penalty: Penalty for repeating tokens.
+            eos_token_id: Token ID that signals end of generation.
+
+        Returns:
+            Generated token IDs including the input prefix.
+        """
+        if eos_token_id is None:
+            eos_token_id = self.config.eos_token_id
+
+        self.eval()
+        generated = input_ids.clone()
+
+        for _ in range(max_new_tokens):
+            # Truncate to max_seq_len if needed
+            context = generated[:, -self.config.max_seq_len:]
+
+            # Forward pass
+            output = self.forward(context)
+            next_token_logits = output["logits"][:, -1, :]  # (batch_size, vocab_size)
+
+            # Apply repetition penalty
+            if repetition_penalty != 1.0:
+                for token_id in set(generated[0].tolist()):
+                    next_token_logits[0, token_id] /= repetition_penalty
+
+            # Apply temperature
+            next_token_logits = next_token_logits / temperature
+
+            # Top-K filtering
+            if top_k > 0:
+                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                next_token_logits[indices_to_remove] = float("-inf")
+
+            # Top-P (nucleus) filtering
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+                # Remove tokens with cumulative probability above threshold
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                indices_to_remove = sorted_indices_to_remove.scatter(
+                    1, sorted_indices, sorted_indices_to_remove
+                )
+                next_token_logits[indices_to_remove] = float("-inf")
+
+            # Sample from the distribution
+            probs = F.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+            # Append generated token
+            generated = torch.cat([generated, next_token], dim=1)
+
+            # Stop if EOS token is generated
+            if next_token.item() == eos_token_id:
+                break
+
+        return generated
+
+    def count_parameters(self) -> int:
+        """Count total trainable parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    def get_parameter_summary(self) -> str:
+        """Get a formatted summary of model parameters."""
+        total = self.count_parameters()
+        lines = [
+            f"{'='*50}",
+            f"  ViuAI-1M Parameter Summary",
+            f"{'='*50}",
+        ]
+        for name, param in self.named_parameters():
+            lines.append(f"  {name:45s} {param.numel():>10,}")
+        lines.append(f"{'='*50}")
+        lines.append(f"  {'Total Parameters':45s} {total:>10,}")
+        lines.append(f"  {'Approx Size (MB)':45s} {total * 4 / 1024 / 1024:>10.2f}")
+        lines.append(f"{'='*50}")
+        return "\n".join(lines)

viuai/tokenizer/__init__.py

@@ -0,0 +1,3 @@
+from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
+
+__all__ = ["ViuAITokenizer"]

viuai/tokenizer/bpe_tokenizer.py

@@ -0,0 +1,349 @@
+"""
+ViuAI-1M BPE Tokenizer — Built from Scratch.
+
+Implements Byte-Pair Encoding (BPE) tokenization:
+1. Start with individual characters as the initial vocabulary
+2. Iteratively merge the most frequent adjacent pair of tokens
+3. Continue until desired vocabulary size is reached
+
+The tokenizer also handles special tokens for chat formatting:
+    <pad>  — padding token
+    <bos>  — beginning of sequence
+    <eos>  — end of sequence
+    <unk>  — unknown token
+    <|user|>    — start of user message
+    <|assistant|> — start of assistant message
+    <|system|>  — start of system message
+"""
+
+import json
+import re
+import os
+from collections import Counter
+from typing import List, Dict, Tuple, Optional
+
+
+# Special tokens with reserved IDs
+SPECIAL_TOKENS = {
+    "<pad>": 0,
+    "<bos>": 1,
+    "<eos>": 2,
+    "<unk>": 3,
+    "<|user|>": 4,
+    "<|assistant|>": 5,
+    "<|system|>": 6,
+}
+
+# Regex pattern to identify special tokens in text
+SPECIAL_TOKEN_PATTERN = re.compile(
+    r"(<pad>|<bos>|<eos>|<unk>|<\|user\|>|<\|assistant\|>|<\|system\|>)"
+)
+
+
+class ViuAITokenizer:
+    """Byte-Pair Encoding tokenizer for the ViuAI-1M model.
+
+    This tokenizer is trained from scratch on the training corpus.
+    It learns subword units by iteratively merging frequent character pairs.
+    """
+
+    def __init__(self, vocab_size: int = 4096):
+        self.target_vocab_size = vocab_size
+
+        # Initialize vocabulary with special tokens
+        self.token_to_id: Dict[str, int] = dict(SPECIAL_TOKENS)
+        self.id_to_token: Dict[int, str] = {v: k for k, v in SPECIAL_TOKENS.items()}
+
+        # BPE merge rules: list of (token_a, token_b) pairs in merge order
+        self.merges: List[Tuple[str, str]] = []
+
+        # Pre-tokenization pattern: split on whitespace boundaries and punctuation
+        self.pre_tokenize_pattern = re.compile(
+            r"""'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?\d+| ?[^\s\w]+|\s+"""
+        )
+
+    def train(self, texts: List[str], verbose: bool = True) -> None:
+        """Train the BPE tokenizer on a corpus of texts.
+
+        Args:
+            texts: List of text strings to train on.
+            verbose: Whether to print progress.
+        """
+        if verbose:
+            print(f"[ViuAI Tokenizer] Training BPE tokenizer...")
+            print(f"  Target vocab size: {self.target_vocab_size}")
+            print(f"  Training on {len(texts)} texts")
+
+        # Step 1: Pre-tokenize and count word frequencies
+        word_freqs = Counter()
+        for text in texts:
+            # Split text into words/chunks
+            words = self.pre_tokenize_pattern.findall(text.lower())
+            for word in words:
+                # Represent each word as a tuple of characters
+                word_freqs[tuple(word)] += 1
+
+        if verbose:
+            print(f"  Unique words: {len(word_freqs)}")
+
+        # Step 2: Initialize vocabulary with all individual characters
+        char_vocab = set()
+        for word in word_freqs:
+            for char in word:
+                char_vocab.add(char)
+
+        # Add characters to vocabulary (after special tokens)
+        next_id = len(self.token_to_id)
+        for char in sorted(char_vocab):
+            if char not in self.token_to_id:
+                self.token_to_id[char] = next_id
+                self.id_to_token[next_id] = char
+                next_id += 1
+
+        if verbose:
+            print(f"  Initial vocab (chars): {len(self.token_to_id)}")
+
+        # Step 3: Iteratively merge most frequent pairs
+        # word_splits tracks how each word is currently split
+        word_splits = {word: list(word) for word in word_freqs}
+
+        num_merges = self.target_vocab_size - len(self.token_to_id)
+        for merge_idx in range(num_merges):
+            # Count all adjacent pairs across the corpus
+            pair_freqs = Counter()
+            for word, freq in word_freqs.items():
+                splits = word_splits[word]
+                for i in range(len(splits) - 1):
+                    pair = (splits[i], splits[i + 1])
+                    pair_freqs[pair] += freq
+
+            if not pair_freqs:
+                if verbose:
+                    print(f"  No more pairs to merge at step {merge_idx}")
+                break
+
+            # Find the most frequent pair
+            best_pair = pair_freqs.most_common(1)[0][0]
+            best_freq = pair_freqs[best_pair]
+
+            # Create new merged token
+            new_token = best_pair[0] + best_pair[1]
+            self.token_to_id[new_token] = next_id
+            self.id_to_token[next_id] = new_token
+            self.merges.append(best_pair)
+            next_id += 1
+
+            # Apply the merge to all words
+            for word in word_splits:
+                splits = word_splits[word]
+                new_splits = []
+                i = 0
+                while i < len(splits):
+                    if (
+                        i < len(splits) - 1
+                        and splits[i] == best_pair[0]
+                        and splits[i + 1] == best_pair[1]
+                    ):
+                        new_splits.append(new_token)
+                        i += 2
+                    else:
+                        new_splits.append(splits[i])
+                        i += 1
+                word_splits[word] = new_splits
+
+            if verbose and (merge_idx + 1) % 500 == 0:
+                print(
+                    f"  Merge {merge_idx + 1}/{num_merges}: "
+                    f"'{best_pair[0]}' + '{best_pair[1]}' → '{new_token}' "
+                    f"(freq: {best_freq})"
+                )
+
+        if verbose:
+            print(f"  Final vocab size: {len(self.token_to_id)}")
+            print(f"  Total merges: {len(self.merges)}")
+            print(f"[ViuAI Tokenizer] Training complete!")
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
+        """Encode text to token IDs.
+
+        Args:
+            text: Input text string.
+            add_special_tokens: Whether to add BOS/EOS tokens.
+
+        Returns:
+            List of token IDs.
+        """
+        token_ids = []
+
+        if add_special_tokens:
+            token_ids.append(SPECIAL_TOKENS["<bos>"])
+
+        # Handle special tokens in the text
+        parts = SPECIAL_TOKEN_PATTERN.split(text)
+
+        for part in parts:
+            if part in SPECIAL_TOKENS:
+                token_ids.append(SPECIAL_TOKENS[part])
+            elif part:
+                # Pre-tokenize into words
+                words = self.pre_tokenize_pattern.findall(part.lower())
+                for word in words:
+                    # Apply BPE merges
+                    word_tokens = self._apply_bpe(word)
+                    for token in word_tokens:
+                        token_ids.append(
+                            self.token_to_id.get(token, SPECIAL_TOKENS["<unk>"])
+                        )
+
+        if add_special_tokens:
+            token_ids.append(SPECIAL_TOKENS["<eos>"])
+
+        return token_ids
+
+    def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str:
+        """Decode token IDs back to text.
+
+        Args:
+            token_ids: List of token IDs.
+            skip_special_tokens: Whether to skip special tokens in output.
+
+        Returns:
+            Decoded text string.
+        """
+        tokens = []
+        for tid in token_ids:
+            token = self.id_to_token.get(tid, "<unk>")
+            if skip_special_tokens and token in SPECIAL_TOKENS:
+                continue
+            tokens.append(token)
+        return "".join(tokens)
+
+    def _apply_bpe(self, word: str) -> List[str]:
+        """Apply learned BPE merges to a word.
+
+        Args:
+            word: Input word string.
+
+        Returns:
+            List of BPE tokens.
+        """
+        # Start with individual characters
+        tokens = list(word)
+
+        # Apply merges in order
+        for merge_a, merge_b in self.merges:
+            new_tokens = []
+            i = 0
+            while i < len(tokens):
+                if (
+                    i < len(tokens) - 1
+                    and tokens[i] == merge_a
+                    and tokens[i + 1] == merge_b
+                ):
+                    new_tokens.append(merge_a + merge_b)
+                    i += 2
+                else:
+                    new_tokens.append(tokens[i])
+                    i += 1
+            tokens = new_tokens
+
+        return tokens
+
+    def format_chat(
+        self,
+        messages: List[Dict[str, str]],
+    ) -> str:
+        """Format a list of chat messages into the ViuAI chat template.
+
+        Chat format:
+            <bos><|system|>System message<|user|>User message<|assistant|>Response<eos>
+
+        Args:
+            messages: List of dicts with 'role' and 'content' keys.
+                     Roles: 'system', 'user', 'assistant'.
+
+        Returns:
+            Formatted chat string.
+        """
+        formatted = ""
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            if role == "system":
+                formatted += f"<|system|>{content}"
+            elif role == "user":
+                formatted += f"<|user|>{content}"
+            elif role == "assistant":
+                formatted += f"<|assistant|>{content}"
+        return formatted
+
+    def save(self, directory: str) -> None:
+        """Save tokenizer to a directory.
+
+        Saves:
+            - vocab.json: token-to-ID mapping
+            - merges.txt: BPE merge rules
+        """
+        os.makedirs(directory, exist_ok=True)
+
+        # Save vocabulary
+        vocab_path = os.path.join(directory, "vocab.json")
+        with open(vocab_path, "w", encoding="utf-8") as f:
+            json.dump(self.token_to_id, f, ensure_ascii=False, indent=2)
+
+        # Save merges
+        merges_path = os.path.join(directory, "merges.txt")
+        with open(merges_path, "w", encoding="utf-8") as f:
+            f.write("# ViuAI-1M BPE Merges\n")
+            f.write(f"# Total merges: {len(self.merges)}\n")
+            for a, b in self.merges:
+                f.write(f"{a} {b}\n")
+
+        print(f"[ViuAI Tokenizer] Saved to {directory}")
+
+    @classmethod
+    def load(cls, directory: str) -> "ViuAITokenizer":
+        """Load tokenizer from a directory.
+
+        Args:
+            directory: Path containing vocab.json and merges.txt.
+
+        Returns:
+            Loaded ViuAITokenizer instance.
+        """
+        tokenizer = cls()
+
+        # Load vocabulary
+        vocab_path = os.path.join(directory, "vocab.json")
+        with open(vocab_path, "r", encoding="utf-8") as f:
+            tokenizer.token_to_id = json.load(f)
+            # Convert string keys back to int values
+            tokenizer.token_to_id = {
+                k: int(v) for k, v in tokenizer.token_to_id.items()
+            }
+        tokenizer.id_to_token = {v: k for k, v in tokenizer.token_to_id.items()}
+
+        # Load merges
+        merges_path = os.path.join(directory, "merges.txt")
+        tokenizer.merges = []
+        with open(merges_path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("#"):
+                    parts = line.split(" ", 1)
+                    if len(parts) == 2:
+                        tokenizer.merges.append((parts[0], parts[1]))
+
+        tokenizer.target_vocab_size = len(tokenizer.token_to_id)
+        print(f"[ViuAI Tokenizer] Loaded from {directory}")
+        print(f"  Vocab size: {len(tokenizer.token_to_id)}")
+        print(f"  Merges: {len(tokenizer.merges)}")
+
+        return tokenizer
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.token_to_id)
+
+    def __len__(self) -> int:
+        return self.vocab_size

viuai/training/__init__.py

@@ -0,0 +1,4 @@
+from viuai.training.dataset import ViuAIDataset, ChatDataCollator
+from viuai.training.trainer import ViuAITrainer
+
+__all__ = ["ViuAIDataset", "ChatDataCollator", "ViuAITrainer"]

viuai/training/dataset.py

@@ -0,0 +1,132 @@
+"""
+ViuAI-1M Dataset and Data Collation.
+
+Handles loading conversation data, tokenizing it using the ViuAI chat template,
+and creating PyTorch DataLoader-compatible datasets.
+
+Supported data format (JSONL):
+    {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
+"""
+
+import json
+import os
+import torch
+from torch.utils.data import Dataset
+from typing import List, Dict, Optional
+
+from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
+
+
+class ViuAIDataset(Dataset):
+    """Dataset for chat-formatted conversation data.
+
+    Loads conversations from JSONL files and tokenizes them
+    using the ViuAI chat template.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer: ViuAITokenizer,
+        max_seq_len: int = 512,
+    ):
+        """Initialize the dataset.
+
+        Args:
+            data_path: Path to a JSONL file with chat conversations.
+            tokenizer: Trained ViuAI tokenizer.
+            max_seq_len: Maximum sequence length (truncate longer sequences).
+        """
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+        self.examples: List[Dict] = []
+
+        # Load conversations from JSONL
+        if os.path.exists(data_path):
+            with open(data_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line:
+                        try:
+                            self.examples.append(json.loads(line))
+                        except json.JSONDecodeError:
+                            continue
+
+        print(f"[ViuAI Dataset] Loaded {len(self.examples)} conversations from {data_path}")
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Get a tokenized training example.
+
+        Returns:
+            Dictionary with 'input_ids' and 'labels' tensors.
+            Labels are the same as input_ids (shifted internally by the model).
+        """
+        example = self.examples[idx]
+        messages = example.get("messages", [])
+
+        # Format using chat template
+        formatted = self.tokenizer.format_chat(messages)
+
+        # Tokenize
+        token_ids = self.tokenizer.encode(formatted, add_special_tokens=True)
+
+        # Truncate to max length
+        if len(token_ids) > self.max_seq_len:
+            token_ids = token_ids[: self.max_seq_len]
+
+        input_ids = torch.tensor(token_ids, dtype=torch.long)
+        labels = input_ids.clone()
+
+        return {"input_ids": input_ids, "labels": labels}
+
+
+class ChatDataCollator:
+    """Collates variable-length sequences into a padded batch.
+
+    Pads sequences to the longest sequence in the batch and
+    creates attention-appropriate labels (pad positions = -100).
+    """
+
+    def __init__(self, pad_token_id: int = 0):
+        self.pad_token_id = pad_token_id
+
+    def __call__(self, batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+        """Collate a list of examples into a batch.
+
+        Args:
+            batch: List of dicts, each with 'input_ids' and 'labels'.
+
+        Returns:
+            Padded batch dict with 'input_ids' and 'labels'.
+        """
+        # Find max length in this batch
+        max_len = max(example["input_ids"].size(0) for example in batch)
+
+        input_ids_list = []
+        labels_list = []
+
+        for example in batch:
+            seq_len = example["input_ids"].size(0)
+            padding_len = max_len - seq_len
+
+            # Pad input_ids with pad_token_id
+            padded_input = torch.cat([
+                example["input_ids"],
+                torch.full((padding_len,), self.pad_token_id, dtype=torch.long),
+            ])
+            input_ids_list.append(padded_input)
+
+            # Pad labels with -100 (ignored in loss computation)
+            padded_labels = torch.cat([
+                example["labels"],
+                torch.full((padding_len,), -100, dtype=torch.long),
+            ])
+            labels_list.append(padded_labels)
+
+        return {
+            "input_ids": torch.stack(input_ids_list),
+            "labels": torch.stack(labels_list),
+        }

viuai/training/trainer.py

@@ -0,0 +1,346 @@
+"""
+ViuAI-1M Training Loop.
+
+Implements a custom training loop with:
+- AdamW optimizer with weight decay
+- Cosine learning rate scheduler with warmup
+- Gradient clipping
+- Training metrics logging
+- Checkpoint saving and loading
+"""
+
+import os
+import math
+import time
+import json
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from typing import Optional, Dict
+
+from viuai.model.viuai_model import ViuAIForCausalLM
+from viuai.model.config import ViuAIConfig
+
+
+class ViuAITrainer:
+    """Custom trainer for the ViuAI-1M model.
+
+    Provides a complete training pipeline with learning rate scheduling,
+    gradient clipping, checkpointing, and metrics logging.
+    """
+
+    def __init__(
+        self,
+        model: ViuAIForCausalLM,
+        train_dataloader: DataLoader,
+        eval_dataloader: Optional[DataLoader] = None,
+        learning_rate: float = 3e-4,
+        weight_decay: float = 0.01,
+        warmup_steps: int = 100,
+        max_grad_norm: float = 1.0,
+        num_epochs: int = 10,
+        save_dir: str = "checkpoints",
+        log_interval: int = 10,
+        eval_interval: int = 100,
+        device: Optional[str] = None,
+    ):
+        self.model = model
+        self.train_dataloader = train_dataloader
+        self.eval_dataloader = eval_dataloader
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.warmup_steps = warmup_steps
+        self.max_grad_norm = max_grad_norm
+        self.num_epochs = num_epochs
+        self.save_dir = save_dir
+        self.log_interval = log_interval
+        self.eval_interval = eval_interval
+
+        # Auto-detect device
+        if device:
+            self.device = torch.device(device)
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
+
+        print(f"[ViuAI Trainer] Using device: {self.device}")
+
+        # Move model to device
+        self.model = self.model.to(self.device)
+
+        # Setup optimizer (AdamW with weight decay)
+        self.optimizer = self._create_optimizer()
+
+        # Training state
+        self.global_step = 0
+        self.best_eval_loss = float("inf")
+        self.training_log: list = []
+
+    def _create_optimizer(self) -> torch.optim.Optimizer:
+        """Create AdamW optimizer with separate weight decay groups.
+
+        Weight decay is NOT applied to bias and normalization parameters.
+        """
+        decay_params = []
+        no_decay_params = []
+
+        for name, param in self.model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if "norm" in name or "bias" in name:
+                no_decay_params.append(param)
+            else:
+                decay_params.append(param)
+
+        param_groups = [
+            {"params": decay_params, "weight_decay": self.weight_decay},
+            {"params": no_decay_params, "weight_decay": 0.0},
+        ]
+
+        return torch.optim.AdamW(
+            param_groups,
+            lr=self.learning_rate,
+            betas=(0.9, 0.95),
+            eps=1e-8,
+        )
+
+    def _get_lr(self, step: int, total_steps: int) -> float:
+        """Compute learning rate with linear warmup and cosine decay.
+
+        Args:
+            step: Current training step.
+            total_steps: Total number of training steps.
+
+        Returns:
+            Current learning rate.
+        """
+        if step < self.warmup_steps:
+            # Linear warmup
+            return self.learning_rate * (step / max(1, self.warmup_steps))
+        else:
+            # Cosine decay
+            progress = (step - self.warmup_steps) / max(1, total_steps - self.warmup_steps)
+            return self.learning_rate * 0.5 * (1.0 + math.cos(math.pi * progress))
+
+    def _update_lr(self, step: int, total_steps: int) -> float:
+        """Update the learning rate for the current step."""
+        lr = self._get_lr(step, total_steps)
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = lr
+        return lr
+
+    def train(self) -> Dict:
+        """Run the full training loop.
+
+        Returns:
+            Dictionary with training metrics.
+        """
+        total_steps = len(self.train_dataloader) * self.num_epochs
+        os.makedirs(self.save_dir, exist_ok=True)
+
+        print(f"\n{'='*60}")
+        print(f"  ViuAI-1M Training")
+        print(f"{'='*60}")
+        print(f"  Parameters:     {self.model.count_parameters():,}")
+        print(f"  Epochs:         {self.num_epochs}")
+        print(f"  Total steps:    {total_steps}")
+        print(f"  Learning rate:  {self.learning_rate}")
+        print(f"  Weight decay:   {self.weight_decay}")
+        print(f"  Warmup steps:   {self.warmup_steps}")
+        print(f"  Device:         {self.device}")
+        print(f"{'='*60}\n")
+
+        self.model.train()
+        start_time = time.time()
+
+        for epoch in range(self.num_epochs):
+            epoch_loss = 0.0
+            epoch_steps = 0
+
+            for batch in self.train_dataloader:
+                # Move batch to device
+                input_ids = batch["input_ids"].to(self.device)
+                labels = batch["labels"].to(self.device)
+
+                # Update learning rate
+                current_lr = self._update_lr(self.global_step, total_steps)
+
+                # Forward pass
+                output = self.model(input_ids=input_ids, labels=labels)
+                loss = output["loss"]
+
+                # Backward pass
+                self.optimizer.zero_grad()
+                loss.backward()
+
+                # Gradient clipping
+                grad_norm = torch.nn.utils.clip_grad_norm_(
+                    self.model.parameters(), self.max_grad_norm
+                )
+
+                # Update weights
+                self.optimizer.step()
+
+                # Track metrics
+                epoch_loss += loss.item()
+                epoch_steps += 1
+                self.global_step += 1
+
+                # Log progress
+                if self.global_step % self.log_interval == 0:
+                    avg_loss = epoch_loss / epoch_steps
+                    elapsed = time.time() - start_time
+                    tokens_per_sec = (
+                        self.global_step * input_ids.shape[0] * input_ids.shape[1]
+                    ) / elapsed
+
+                    log_entry = {
+                        "step": self.global_step,
+                        "epoch": epoch + 1,
+                        "loss": loss.item(),
+                        "avg_loss": avg_loss,
+                        "lr": current_lr,
+                        "grad_norm": grad_norm.item(),
+                        "tokens_per_sec": tokens_per_sec,
+                    }
+                    self.training_log.append(log_entry)
+
+                    print(
+                        f"  Step {self.global_step:>6d}/{total_steps} | "
+                        f"Epoch {epoch+1}/{self.num_epochs} | "
+                        f"Loss: {loss.item():.4f} | "
+                        f"Avg Loss: {avg_loss:.4f} | "
+                        f"LR: {current_lr:.2e} | "
+                        f"Grad Norm: {grad_norm:.3f} | "
+                        f"Tok/s: {tokens_per_sec:.0f}"
+                    )
+
+                # Evaluation
+                if self.eval_dataloader and self.global_step % self.eval_interval == 0:
+                    eval_loss = self.evaluate()
+                    if eval_loss < self.best_eval_loss:
+                        self.best_eval_loss = eval_loss
+                        self.save_checkpoint("best")
+                        print(f"  ★ New best eval loss: {eval_loss:.4f}")
+                    self.model.train()
+
+            # End of epoch
+            avg_epoch_loss = epoch_loss / max(1, epoch_steps)
+            print(f"\n  [Epoch {epoch+1}/{self.num_epochs}] Average Loss: {avg_epoch_loss:.4f}\n")
+
+            # Save checkpoint at end of each epoch
+            self.save_checkpoint(f"epoch_{epoch+1}")
+
+        # Training complete
+        total_time = time.time() - start_time
+        print(f"\n{'='*60}")
+        print(f"  Training Complete!")
+        print(f"  Total time: {total_time:.1f}s")
+        print(f"  Best eval loss: {self.best_eval_loss:.4f}")
+        print(f"{'='*60}\n")
+
+        # Save final checkpoint and training log
+        self.save_checkpoint("final")
+        self._save_training_log()
+
+        return {
+            "total_steps": self.global_step,
+            "total_time": total_time,
+            "best_eval_loss": self.best_eval_loss,
+            "final_loss": avg_epoch_loss,
+        }
+
+    @torch.no_grad()
+    def evaluate(self) -> float:
+        """Run evaluation on the eval dataset.
+
+        Returns:
+            Average evaluation loss.
+        """
+        self.model.eval()
+        total_loss = 0.0
+        total_steps = 0
+
+        for batch in self.eval_dataloader:
+            input_ids = batch["input_ids"].to(self.device)
+            labels = batch["labels"].to(self.device)
+
+            output = self.model(input_ids=input_ids, labels=labels)
+            total_loss += output["loss"].item()
+            total_steps += 1
+
+        avg_loss = total_loss / max(1, total_steps)
+        print(f"  [Eval] Loss: {avg_loss:.4f}")
+        return avg_loss
+
+    def save_checkpoint(self, name: str) -> None:
+        """Save model checkpoint.
+
+        Args:
+            name: Checkpoint name (e.g., 'best', 'epoch_1', 'final').
+        """
+        checkpoint_dir = os.path.join(self.save_dir, name)
+        os.makedirs(checkpoint_dir, exist_ok=True)
+
+        # Save model weights
+        torch.save(
+            self.model.state_dict(),
+            os.path.join(checkpoint_dir, "model.pt"),
+        )
+
+        # Save optimizer state
+        torch.save(
+            self.optimizer.state_dict(),
+            os.path.join(checkpoint_dir, "optimizer.pt"),
+        )
+
+        # Save training state
+        state = {
+            "global_step": self.global_step,
+            "best_eval_loss": self.best_eval_loss,
+        }
+        with open(os.path.join(checkpoint_dir, "training_state.json"), "w") as f:
+            json.dump(state, f, indent=2)
+
+        # Save config
+        self.model.config.save(os.path.join(checkpoint_dir, "config.json"))
+
+        print(f"  [Checkpoint] Saved: {checkpoint_dir}")
+
+    def load_checkpoint(self, name: str) -> None:
+        """Load model checkpoint.
+
+        Args:
+            name: Checkpoint name to load.
+        """
+        checkpoint_dir = os.path.join(self.save_dir, name)
+
+        # Load model weights
+        self.model.load_state_dict(
+            torch.load(os.path.join(checkpoint_dir, "model.pt"), map_location=self.device)
+        )
+
+        # Load optimizer state
+        opt_path = os.path.join(checkpoint_dir, "optimizer.pt")
+        if os.path.exists(opt_path):
+            self.optimizer.load_state_dict(
+                torch.load(opt_path, map_location=self.device)
+            )
+
+        # Load training state
+        state_path = os.path.join(checkpoint_dir, "training_state.json")
+        if os.path.exists(state_path):
+            with open(state_path, "r") as f:
+                state = json.load(f)
+            self.global_step = state.get("global_step", 0)
+            self.best_eval_loss = state.get("best_eval_loss", float("inf"))
+
+        print(f"  [Checkpoint] Loaded: {checkpoint_dir}")
+
+    def _save_training_log(self) -> None:
+        """Save training metrics log to JSON."""
+        log_path = os.path.join(self.save_dir, "training_log.json")
+        with open(log_path, "w") as f:
+            json.dump(self.training_log, f, indent=2)
+        print(f"  [Log] Saved training log to {log_path}")