Update README.md

README.md

@@ -1,6 +1,6 @@
 ---
 language: en
-license: mit
+license: apache-2.0
 library_name: pytorch
 tags:
 - transformer
@@ -12,369 +12,123 @@ tags:
 
 # Model Card for LumenBase
 
-<!-- Provide a quick summary of what the model is/does. -->
-
-LumenBase is a 128M parameter GPT-style transformer language model built entirely from scratch for educational and research purposes. The model features modern architectural components including Grouped Multi-Query Attention (GQA), SwiGLU activation, RMSNorm, and Rotary Position Embeddings (RoPE).
+A 128M parameter GPT-style transformer built from scratch for educational purposes, featuring Grouped Multi-Query Attention (GQA), SwiGLU, RMSNorm, and RoPE.
 
 ## Model Details
 
 ### Model Description
 
-<!-- Provide a longer summary of what this model is. -->
-
-LumenBase is a foundational language model created entirely from scratch to explore every step of modern LLM development — from data preprocessing and tokenization to architecture design, training, and evaluation. This project implements a decoder-only transformer architecture with several modern optimizations:
-
-- **Grouped Multi-Query Attention (GQA)**: Efficient attention mechanism with 12 query heads and 4 key-value heads (3 groups)
-- **SwiGLU Activation**: Advanced feed-forward network activation function
-- **RMSNorm**: Layer normalization for improved training stability
-- **Rotary Position Embeddings (RoPE)**: Relative position encoding
-- **Weight Tying**: Shared weights between embedding and output layers
+LumenBase is a decoder-only transformer language model implementing modern architectural optimizations:
+- **Architecture**: 12-layer transformer with GQA (12 query heads, 4 KV heads), SwiGLU activation, RMSNorm, and RoPE
+- **Parameters**: 128M (768 hidden size, 3072 FFN, 2048 context length)
+- **Training**: Mixed precision (FP16/BF16) with custom tokenizer (32K vocab)
 
-The model was trained on custom datasets using mixed precision training (FP16/BF16) with gradient accumulation, cosine annealing scheduler with linear warmup, and automatic checkpointing.
-
-- **Developed by:** Hariom Jangra (HariomJangra)
-- **Model type:** Decoder-only Transformer Language Model
-- **Language(s) (NLP):** English
+- **Developed by:** Hariom Jangra
+- **Model type:** Decoder-only Transformer
+- **Language:** English
 - **License:** MIT
-- **Finetuned from model [optional]:** N/A (trained from scratch)
-
-### Model Sources [optional]
-
-<!-- Provide the basic links for the model. -->
-
 - **Repository:** https://github.com/HariomJangra/project-lumen
-- **Paper [optional]:** N/A
-- **Demo [optional]:** N/A
 
 ## Uses
 
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-
-### Direct Use
+**Direct Use:**
+- Text generation and completion
+- Educational resource for understanding transformer architecture
+- Research baseline for language models
+- Foundation for fine-tuning on specific tasks
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+**Downstream Use:**
+- Instruction tuning
+- Chat applications
+- Domain-specific fine-tuning
 
-LumenBase can be used directly for:
-- Text generation and completion tasks
-- Educational purposes to understand transformer architecture and training
-- Research on language model behavior and capabilities
-- Baseline for fine-tuning on specific downstream tasks
-- Understanding modern LLM architectural components
+**Out-of-Scope:**
+- Production deployments
+- Safety-critical applications
+- Applications requiring factual accuracy without verification
+- This is an educational model - use established frameworks for production
 
-### Downstream Use [optional]
+## Limitations
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+**Technical:**
+- Limited size (128M parameters) - below state-of-the-art performance
+- 2048 token context window
+- May generate incoherent text for complex prompts
 
-The model can be fine-tuned for:
-- Instruction following
-- Chat-based applications
-- Domain-specific text generation
-- Task-specific adaptations
-- Further research on specialized applications
+**Bias & Safety:**
+- May perpetuate training data biases
+- Not evaluated for fairness across demographics
+- Can generate inappropriate content
+- Should not be relied upon for factual information
 
-### Out-of-Scope Use
+**Recommendations:** This is an educational model. Verify all outputs, implement content filtering for applications, and use production-ready models for commercial use.
 
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+## Training
 
-This model is **not suitable** for:
-- Production deployments requiring high-quality generation
-- Safety-critical applications
-- Applications requiring factual accuracy without verification
-- Generation of harmful, hateful, or biased content
-- Large-scale commercial applications without proper evaluation
+**Data:** Custom datasets tokenized with BPE (32K vocab)
 
-This is an educational/research implementation. For production use, consider established frameworks like Hugging Face Transformers.
+**Hyperparameters:**
+- Optimizer: AdamW (lr=3e-4, weight_decay=0.1)
+- Batch: 12 × 4 gradient accumulation = 48 effective
+- Sequence length: 2048 tokens
+- Scheduler: Linear warmup + Cosine annealing
+- Precision: Mixed (FP16/BF16/FP32)
+- Dropout: 0.1 (training), 0.0 (inference)
 
-## Bias, Risks, and Limitations
+![Training Loss](PreTraining/images/training_loss_curve.png)
 
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+## Evaluation
 
-**Technical Limitations:**
-- Limited model size (128M parameters) compared to larger production models
-- Performance on benchmarks is below state-of-the-art models
-- May generate incoherent or nonsensical text for complex prompts
-- Limited context window (2048 tokens)
+Evaluated on standard NLP benchmarks:
 
-**Bias and Social Limitations:**
-- The model may perpetuate biases present in training data
-- Not evaluated for fairness across different demographic groups
-- May generate inappropriate or offensive content
-- Should not be relied upon for factual information without verification
+| Benchmark | Accuracy | Correct/Total |
+|-----------|----------|---------------|
+| **ARC-Easy** | 39.48% | 938/2,376 |
+| **ARC-Challenge** | 23.55% | 276/1,172 |
+| **HellaSwag** | 32.62% | 334/1,024 |
 
-**Research/Educational Nature:**
-- This is a learning project, not optimized for production use
-- Training data sources and quality may vary
-- Limited testing across diverse use cases
+**Summary:** Baseline performance consistent with a 128M educational model. Results show capability on easier tasks with room for improvement on complex reasoning.
 
-### Recommendations
+## Technical Specifications
 
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+**Architecture:** Decoder-only Transformer
+- 12 layers, 768 hidden size, 12 attention heads (4 KV heads)
+- SwiGLU FFN (3072 intermediate), RMSNorm, RoPE
+- 32K vocab, 2048 max sequence length
+- Weight tying between embedding and output layers
 
-Users (both direct and downstream) should:
-- Be aware this is an educational model with limited capabilities
-- Not use for safety-critical or production applications
-- Verify all generated content before use
-- Implement appropriate content filtering for downstream applications
-- Consider the model's limitations when interpreting results
-- Use established production-ready models for commercial applications
+**Implementation:** Custom PyTorch implementation from scratch
 
-## How to Get Started with the Model
+**Software:** Python 3.13, PyTorch, NumPy, Tokenizers, tqdm, matplotlib
 
-Use the code below to get started with the model.
+## How to Use
 
 ```python
 import torch
 from ModelArchitecture import Transformer, ModelConfig, generate
 from tokenizers import Tokenizer
 
-# Load model configuration
-config = ModelConfig(
-    vocab_size=32000,
-    hidden_size=768,
-    n_heads=12,
-    n_kv_heads=4,
-    n_kv_groups=3,
-    head_dim=64,
-    n_layers=12,
-    attention_bias=False,
-    intermediate_size=3072,
-    mlp_bias=False,
-    eps=1e-5,
-    dropout=0.0,
-    max_position_embeddings=2048,
-    pre_norm=True,
-    tie_weights=True,
-    max_seq_len=2048
-)
-
-# Initialize model
-model = Transformer(config)
+# Load configuration and model
+config = ModelConfig(vocab_size=32000, hidden_size=768, n_heads=12, 
+                     n_kv_heads=4, n_kv_groups=3, head_dim=64, n_layers=12,
+                     intermediate_size=3072, max_position_embeddings=2048,
+                     dropout=0.0, pre_norm=True, tie_weights=True)
 
-# Load trained weights
-checkpoint = torch.load('LumenBase.safetensors', map_location='cpu')
-model.load_state_dict(checkpoint)
+model = Transformer(config)
+model.load_state_dict(torch.load('LumenBase.safetensors'))
 model.eval()
 
-# Load tokenizer
-tokenizer = Tokenizer.from_file('LumenTokenizer.json')
-
 # Generate text
+tokenizer = Tokenizer.from_file('LumenTokenizer.json')
 prompt = "Once upon a time"
 input_ids = torch.tensor([tokenizer.encode(prompt).ids])
 
-output = generate(
-    model=model,
-    input_ids=input_ids,
-    max_new_tokens=100,
-    temperature=0.8,
-    top_k=50,
-    top_p=0.9,
-    do_sample=True
-)
-
-generated_text = tokenizer.decode(output[0].tolist())
-print(generated_text)
-```
-
-## Training Details
-
-### Training Data
-
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
-
-The model was trained on custom datasets prepared specifically for this project. The training pipeline included:
-
-1. **Dataset Preparation**: Text data collection and preprocessing
-2. **Tokenization**: Custom BPE (Byte Pair Encoding) tokenizer trained with a vocabulary size of 32,000 tokens
-3. **Data Processing**: Text tokenization and conversion to token IDs stored as NumPy arrays
-   - TokenizedDataSet1.npy
-   - TokenizedDataSet2.npy
-   - TokenizedDataset3.npy
-
-The training data was tokenized using a custom-trained tokenizer (LumenTokenizer) optimized for the target domain.
-
-### Training Procedure
-
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-
-#### Preprocessing [optional]
-
-- Text cleaning and normalization
-- BPE tokenization with 32K vocabulary
-- Sequence chunking to 2048 token context windows
-- Data stored in efficient NumPy format for fast loading
-
-#### Training Hyperparameters
-
-- **Optimizer**: AdamW (lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
-- **Scheduler**: Linear warmup (2000 steps) + Cosine annealing
-- **Batch Size**: 12 sequences per batch
-- **Gradient Accumulation Steps**: 4 (effective batch size: 48)
-- **Sequence Length**: 2048 tokens
-- **Dropout**: 0.1 during training, 0.0 during inference
-- **Gradient Clipping**: Max norm 1.0
-- **Training regime**: Mixed precision (automatic FP16/BF16/FP32 based on hardware)
-
-#### Speeds, Sizes, Times [optional]
-
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-
-- **Model Parameters**: 128M (128 million)
-- **Model Size**: ~512 MB (FP32), ~256 MB (FP16)
-- **Checkpoint Frequency**: Every N steps with automatic best model saving
-- **Training monitored with**: Training and validation loss curves
-- **Final checkpoint**: best_model_params_110k.pt
-
-![Training Loss Curve](PreTraining/images/training_loss_curve.png)
-
-## Evaluation
-
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-The model was evaluated on three standard NLP benchmarks:
-
-1. **ARC-Easy** (AI2 Reasoning Challenge - Easy): 2,376 questions
-2. **ARC-Challenge** (AI2 Reasoning Challenge - Challenge): 1,172 questions  
-3. **HellaSwag**: 1,024 examples for commonsense reasoning
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-The model was evaluated on:
-- Multiple-choice question answering
-- Commonsense reasoning
-- Scientific reasoning
-- Reading comprehension
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-**Accuracy**: The primary metric used for all benchmarks, measuring the percentage of correctly answered questions.
-
-### Results
-
-| Benchmark | Accuracy | Correct | Total |
-|-----------|----------|---------|-------|
-| **ARC-Easy** | 39.48% | 938 | 2,376 |
-| **ARC-Challenge** | 23.55% | 276 | 1,172 |
-| **HellaSwag** | 32.62% | 334 | 1,024 |
-
-#### Summary
-
-The LumenBase model demonstrates baseline performance on standard NLP benchmarks. As expected for a 128M parameter model trained from scratch for educational purposes:
-
-- **ARC-Easy**: Achieves ~39% accuracy, showing some capability on easier scientific reasoning tasks
-- **ARC-Challenge**: Scores ~24% on the more difficult version, indicating room for improvement on complex reasoning
-- **HellaSwag**: Reaches ~33% on commonsense reasoning, slightly above random chance (25% for 4-choice questions)
-
-These results are consistent with a small-scale educational model and provide a baseline for future improvements through:
-- Additional training data
-- Longer training duration
-- Model scaling
-- Fine-tuning on specific tasks
-- Improved training techniques
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-**Architecture Details:**
-- **Attention Mechanism**: Grouped Multi-Query Attention reduces KV cache size while maintaining performance
-- **Activation Function**: SwiGLU provides better gradient flow compared to traditional ReLU
-- **Normalization**: RMSNorm (Root Mean Square Layer Normalization) for improved stability
-- **Position Encoding**: RoPE (Rotary Position Embeddings) for better handling of relative positions
-- **Weight Tying**: Embedding and output layer share weights, reducing parameter count
-
-**Key Design Choices:**
-- Decoder-only architecture following GPT design principles
-- Pre-normalization for better training stability
-- Efficient attention with 12 query heads, 4 KV heads (grouped into 3)
-- Intermediate FFN size of 3072 (4x hidden size)
-
-**Implementation Highlights:**
-- Custom implementation from scratch using PyTorch
-- Supports various sampling strategies: greedy, top-k, top-p (nucleus), temperature scaling
-- Gradient accumulation for effective larger batch sizes
-- Automatic mixed precision training support
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** Consumer-grade GPU (specific hardware varies)
-- **Hours used:** Educational project, training time not formally tracked
-- **Cloud Provider:** N/A (local training)
-- **Compute Region:** N/A
-- **Carbon Emitted:** Not formally measured
-
-**Note**: As an educational project, formal carbon footprint tracking was not implemented. Future iterations could benefit from tracking environmental impact.
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-**Architecture**: Decoder-only Transformer (GPT-style)
-
-**Configuration:**
-```python
-vocab_size: 32000
-hidden_size: 768
-n_heads: 12
-n_kv_heads: 4
-n_kv_groups: 3
-head_dim: 64
-n_layers: 12
-intermediate_size: 3072
-max_position_embeddings: 2048
-dropout: 0.1 (training) / 0.0 (inference)
-```
-
-**Key Components:**
-- **Grouped Multi-Query Attention**: 12 query heads, 4 key-value heads
-- **Feed-Forward Network**: SwiGLU activation with 3072 intermediate dimensions
-- **Layer Normalization**: RMSNorm (epsilon=1e-5)
-- **Position Encoding**: Rotary Position Embeddings (RoPE)
-- **Weight Tying**: Shared embedding and output projection weights
-
-**Training Objective**: Causal language modeling with cross-entropy loss
-
-### Compute Infrastructure
-
-Educational project trained on consumer hardware.
-
-#### Hardware
-
-- Consumer-grade GPU (specific configuration varies)
-- Training performed locally, not on cloud infrastructure
-
-#### Software
-
-```
-Python 3.13
-PyTorch (latest)
-NumPy
-Tokenizers (Hugging Face)
-tqdm (progress tracking)
-matplotlib (visualization)
+output = generate(model, input_ids, max_new_tokens=100, 
+                 temperature=0.8, top_k=50, top_p=0.9)
+print(tokenizer.decode(output[0].tolist()))
 ```
 
-**Custom Implementation**: All model components implemented from scratch in PyTorch without using high-level transformer libraries.
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
+## Citation
 
 ```bibtex
 @misc{lumenbase2024,
@@ -382,66 +136,12 @@ matplotlib (visualization)
   title = {LumenBase: A 128M Parameter Language Model Built from Scratch},
   year = {2024},
   publisher = {GitHub},
-  journal = {GitHub repository},
   howpublished = {\url{https://github.com/HariomJangra/project-lumen}}
 }
 ```
 
-**APA:**
-
-Jangra, H. (2024). *LumenBase: A 128M Parameter Language Model Built from Scratch* [Computer software]. GitHub. https://github.com/HariomJangra/project-lumen
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-**Terms:**
-
-- **BPE (Byte Pair Encoding)**: Tokenization algorithm that builds vocabulary by iteratively merging frequent character pairs
-- **GQA (Grouped Multi-Query Attention)**: Attention mechanism where multiple query heads share fewer key-value heads, reducing memory and computation
-- **RMSNorm**: Root Mean Square Layer Normalization - simplified normalization that only rescales using RMS statistics
-- **RoPE (Rotary Position Embeddings)**: Position encoding that encodes absolute positions with rotation matrices and naturally incorporates relative position information
-- **SwiGLU**: Activation function combining Swish activation with Gated Linear Units for improved model performance
-- **Weight Tying**: Technique where embedding and output layers share parameters to reduce model size
-
-**Sampling Strategies:**
-- **Greedy Decoding**: Always select the token with highest probability
-- **Top-k Sampling**: Sample from the k most likely tokens
-- **Top-p (Nucleus) Sampling**: Sample from smallest set of tokens whose cumulative probability exceeds p
-- **Temperature Scaling**: Adjust probability distribution sharpness (lower = more deterministic, higher = more random)
-
-## More Information [optional]
-
-**Project Structure:**
-- `PreTraining/Implementation/`: Training scripts and data preparation notebooks
-- `PreTraining/Benchmark/`: Evaluation scripts and results
-- `PreTraining/Inference/`: Text generation and inference code
-- `PreTraining/Models/`: Saved model checkpoints
-- `PreTraining/DataSets/`: Tokenized training data
-
-**Future Work:**
-- Fine-tuning for instruction following
-- Chat model adaptation
-- Task-specific fine-tuning
-- Scaling to larger model sizes
-- Improved training data curation
-- Advanced sampling techniques
-
-**Learning Resources:**
-This project serves as a comprehensive educational resource covering:
-1. Dataset preparation and cleaning
-2. Custom tokenizer training
-3. Transformer architecture implementation
-4. Training loop with modern optimizations
-5. Evaluation on standard benchmarks
-6. Text generation with various sampling strategies
-
-For detailed implementation and usage, please refer to the [GitHub repository](https://github.com/HariomJangra/project-lumen).
-
-## Model Card Authors [optional]
-
-Hariom Jangra ([@HariomJangra](https://github.com/HariomJangra))
+## Contact
 
-## Model Card Contact
+**Author:** Hariom Jangra ([@HariomJangra](https://github.com/HariomJangra))
 
-For questions or feedback about this model, please open an issue on the [GitHub repository](https://github.com/HariomJangra/project-lumen).
+For questions or feedback, please open an issue on the [GitHub repository](https://github.com/HariomJangra/project-lumen).