Add StoryKimi ZeroGPU implementation

- Add ZeroGPU-compatible app.py with @spaces.GPU decorator
- Copy all necessary model files (config.py, model.py, tokenizer.py, inference.py)
- Update requirements.txt with spaces and gradio dependencies
- Comprehensive README.md based on original StoryKimi with HF Spaces adaptations
- Add .gitignore to exclude checkpoints and temporary files while keeping main model
- Configure metadata for ZeroGPU hardware in README frontmatter

.gitignore

@@ -0,0 +1,217 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be added to the global gitignore or merged into this project gitignore.  For a PyCharm
+#  project, uncomment the following line:
+#.idea/
+
+# Model checkpoints and weights (except the main one we want to keep)
+checkpoints/
+*.pt
+*.pth
+*.ckpt
+*.safetensors
+!checkpoint_2000.pt
+
+# Wandb logs
+wandb/
+runs/
+
+# Generated data
+generated_data/
+data/
+datasets/
+
+# Images (except for README)
+images/
+*.png
+*.jpg
+*.jpeg
+*.gif
+!images/image.png
+
+# Gradio temporary files
+gradio_cached_examples/
+flagged/
+
+# OS files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Temporary files
+*.tmp
+*.temp
+temp/
+
+# Log files
+*.log
+logs/
+
+# Test files
+test_outputs/
+test_results/

README.md

@@ -1,13 +1,148 @@
 ---
 title: StoryKimi Zero
-emoji: 📈
-colorFrom: gray
-colorTo: pink
+emoji: 🚀
+colorFrom: blue
+colorTo: purple
 sdk: gradio
 sdk_version: 5.42.0
 app_file: app.py
 pinned: false
-license: apache-2.0
+license: mit
+hardware: zero-gpu
+short_description: Generate stories with StoryKimi model using ZeroGPU
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# StoryKimi Zero - DeepSeek V3 Inspired Model on ZeroGPU
+
+A PyTorch implementation of a DeepSeek V3 inspired transformer model with Mixture of Experts (MoE), Latent Attention, and other advanced features, deployed on Hugging Face Spaces with ZeroGPU for efficient inference.
+
+![StoryKimi Model](https://huggingface.co/YuvrajSingh9886/StoryKimi/resolve/main/images/image.png)
+
+## 📊 Training Results & Model Weights
+
+**📈 View Training Report**: [StoryKimi Training Results on WandB](https://wandb.ai/rentio/DSV-Training/reports/SmolKimi-A-smaller-Kimi-K2---VmlldzoxMzYwNDQ4Mg?accessToken=lfs6n1y7gn8q0f0dwilta8yuwzxel45ztzbbcavwbqp7jsyv1p7cz9elflycv9fg)
+
+**💾 Pre-trained Weights**: 
+- **Hugging Face Model**: [YuvrajSingh9886/StoryKimi](https://huggingface.co/YuvrajSingh9886/StoryKimi)
+- **WandB Checkpoints**: Check the WandB report above for additional trained model checkpoints
+
+## 🌟 Features
+
+- **ZeroGPU Integration**: Dynamic GPU allocation with NVIDIA H200 slices (70GB VRAM)
+- **Latent Attention**: Efficient attention mechanism with compressed key-value representations
+- **Mixture of Experts (MoE)**: 8 experts with top-2 routing and shared expert support
+- **SWiGLU Activation**: Advanced activation function in expert layers
+- **Sinusoidal Positional Embeddings**: Position encoding for sequence understanding
+- **Interactive Interface**: User-friendly Gradio interface with real-time generation
+- **Multiple Sampling Methods**: Top-k sampling with temperature control
+- **Real-time Generation**: Fast inference with automatic scaling
+
+## 🔧 Model Architecture
+
+### Default Configuration
+- **Embedding Dimensions**: 384
+- **Decoder Layers**: 6
+- **Attention Heads**: 8
+- **MoE Experts**: 8 (top-2 routing)
+- **Block Size**: 128 tokens
+- **Vocabulary Size**: Based on Llama-2-7b tokenizer (~32,000 tokens)
+- **Latent Dimension**: 64 (for compressed attention)
+
+### ZeroGPU Configuration
+- **GPU Type**: NVIDIA H200 slice
+- **Available VRAM**: 70GB per workload
+- **Max Duration**: 120 seconds per generation
+- **Deployment**: Hugging Face Spaces with automatic scaling
+
+## 🎯 Usage
+
+1. **Enter your story prompt** in the text box
+2. **Select model checkpoint** (Checkpoint 2000 available)
+3. **Adjust generation parameters**:
+   - **Max Length**: 10-128 tokens
+   - **Temperature**: 0.1-2.0 (creativity vs coherence)
+   - **Top-k**: 1-100 (vocabulary filtering)
+4. **Click "Generate Text"** to create your AI-generated story
+5. **Enjoy your personalized story!**
+
+## 💡 Generation Tips
+
+- **Lower temperature** (0.1-0.7) for more coherent and focused stories
+- **Higher temperature** (0.8-2.0) for more creative and diverse outputs
+- **Adjust top-k** to control vocabulary diversity and randomness
+- **Use descriptive prompts** for better and more relevant results
+- **Experiment with different lengths** to find your preferred story format
+
+## 🔄 ZeroGPU Benefits
+
+- **Free GPU Access**: No cost for users to generate stories
+- **Efficient Resource Usage**: GPU allocated only when needed for inference
+- **Automatic Scaling**: Handles multiple concurrent users seamlessly
+- **High Performance**: NVIDIA H200 acceleration for fast generation
+- **No Setup Required**: Ready-to-use interface with pre-loaded model
+
+## 🏗️ Technical Implementation
+
+### Model Features
+- **Latent Attention**: Compressed key-value representations for efficiency
+- **Mixture of Experts**: 8 experts with intelligent routing
+- **Advanced Activation**: SWiGLU for better performance
+- **Positional Encoding**: Sinusoidal embeddings for sequence understanding
+
+### Deployment Features
+- **ZeroGPU Decorator**: `@spaces.GPU(duration=120)` for dynamic allocation
+- **Optimized Loading**: Efficient model loading and initialization
+- **Error Handling**: Robust error management for better user experience
+- **Real-time Feedback**: Live generation status and results
+
+## 🚀 Local Development
+
+Want to run this locally or contribute? Check out the full repository:
+
+**📁 Source Code**: [YuvrajSingh-mist/SmolHub/StoryKimi](https://github.com/YuvrajSingh-mist/SmolHub/tree/main/StoryKimi)
+
+### Quick Local Setup
+```bash
+# Clone the repository
+git clone https://github.com/YuvrajSingh-mist/SmolHub.git
+cd SmolHub/StoryKimi
+
+# Install dependencies
+chmod +x install.sh
+./install.sh
+
+# Run Gradio interface
+cd gradio
+python app.py
+```
+
+### Training Your Own Model
+```bash
+# Set your HF token for Llama-2 tokenizer access
+export HF_TOKEN="your_token_here"
+
+# Basic training
+python trainer.py
+
+# Advanced training with custom parameters
+python trainer.py --embeddings_dims 512 --experts 16 --epochs 5
+```
+
+## 📊 Model Performance
+
+The model has been trained on diverse text data and shows strong performance in:
+- **Story Generation**: Creative and coherent narrative creation
+- **Text Continuation**: Natural extension of given prompts
+- **Style Adaptation**: Adapting to different writing styles and genres
+- **Character Development**: Creating consistent characters and dialogue
+
+## 🔗 Related Links
+
+- **Full Project**: [SmolHub Repository](https://github.com/YuvrajSingh-mist/SmolHub)
+- **Model Weights**: [HuggingFace Model](https://huggingface.co/YuvrajSingh9886/StoryKimi)
+- **Training Report**: [WandB Results](https://wandb.ai/rentio/DSV-Training/reports/SmolKimi-A-smaller-Kimi-K2---VmlldzoxMzYwNDQ4Mg?accessToken=lfs6n1y7gn8q0f0dwilta8yuwzxel45ztzbbcavwbqp7jsyv1p7cz9elflycv9fg)
+- **Other Models**: [SmolMixtral](https://github.com/YuvrajSingh-mist/SmolHub/tree/main/SmolMixtral), [SmolTransformer](https://github.com/YuvrajSingh-mist/SmolHub/tree/main/SmolTransformer)
+
+## 📝 License
+
+MIT License - See LICENSE file for details

app.py

@@ -0,0 +1,202 @@
+import gradio as gr
+import spaces  # HF Spaces ZeroGPU decorator - only available in HF Spaces environment
+import torch
+import torch.nn.functional as F
+import os
+import sys
+
+from config import ModelArgs, get_args
+from model import DeepSeekV3, initialize_tokenizer
+from tokenizer import Tokenizer
+from inference import topk_sampling
+
+# Global variables
+tk = None
+model = None
+model_args = None
+
+# Model paths - using the checkpoint in the HF Space
+model_paths = {
+    "Checkpoint 2000": "./checkpoint_2000.pt",
+}
+
+def initialize_app():
+    """Initialize the app with tokenizer and model args"""
+    global tk, model_args
+    
+    # Initialize model args
+    model_args = ModelArgs()
+    
+    # Initialize tokenizer (no HF token needed for basic operation)
+    if tk is None:
+        tk = Tokenizer(hf_token=None)
+        tk = tk.ready_tokenizer()
+    
+    # Initialize the global tokenizer in model.py
+    initialize_tokenizer(hf_token=None)
+
+def load_model(model_path, device, model_args):
+    """Load model from checkpoint"""
+    model = DeepSeekV3(
+        embeddings_dims=model_args.embeddings_dims,
+        block_size=model_args.block_size,
+        vocab_size=model_args.vocab_size,
+        dropout=model_args.dropout,
+        device=device
+    )
+    
+    if os.path.exists(model_path):
+        checkpoint = torch.load(model_path, map_location=device)
+        model.load_state_dict(checkpoint)
+        model.eval()
+        print(f"Model loaded from {model_path}")
+    else:
+        print(f"Checkpoint {model_path} not found. Using randomly initialized model.")
+    
+    return model
+
+@spaces.GPU(duration=120)
+def generate_text(prompt, model_choice, max_length, temperature, top_k):
+    """Generate text using the selected model and top-k sampling"""
+    global tk, model_args
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+    
+    # Load the selected model
+    model_path = model_paths.get(model_choice, "./checkpoint_2000.pt")
+    model = load_model(model_path, device, model_args)
+    model = model.to(device)
+    
+    try:
+        generated_text = topk_sampling(
+            model=model,
+            prompt=prompt,
+            device=device,
+            max_length=max_length,
+            top_k=top_k,
+            temperature=temperature,
+            tokenizer=tk
+        )
+            
+        return generated_text
+        
+    except Exception as e:
+        return f"Error generating text: {str(e)}"
+
+def create_interface():
+    """Create the Gradio interface"""
+    global tk, model_args
+    
+    # Initialize the app
+    initialize_app()
+    
+    with gr.Blocks(title="StoryKimi Text Generator", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🚀 StoryKimi Text Generator")
+        gr.Markdown("Generate text using the Kimi K2 inspired StoryKimi model with ZeroGPU support.")
+        gr.Markdown("⚡ **Powered by ZeroGPU** - Dynamic GPU allocation for efficient inference")
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                prompt_input = gr.Textbox(
+                    label="Input Prompt",
+                    placeholder="Enter your prompt here...",
+                    lines=3,
+                    value="Once upon a time there lived a baby deer named Bambi."
+                )
+                
+                with gr.Row():
+                    model_dropdown = gr.Dropdown(
+                        choices=list(model_paths.keys()),
+                        label="Model Checkpoint",
+                        value="Checkpoint 2000"
+                    )
+                
+                with gr.Row():
+                    max_length_slider = gr.Slider(
+                        minimum=10,
+                        maximum=128,
+                        value=50,
+                        step=10,
+                        label="Max Length"
+                    )
+                    
+                    temperature_slider = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=0.9,
+                        step=0.1,
+                        label="Temperature"
+                    )
+                
+                with gr.Row():
+                    top_k_slider = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                        label="Top-k"
+                    )
+                
+                with gr.Row():
+                    top_k_slider = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                        label="Top-k"
+                    )
+                
+                generate_btn = gr.Button("🎯 Generate Text", variant="primary", size="lg")
+                
+            with gr.Column(scale=3):
+                output_text = gr.Textbox(
+                    label="Generated Text",
+                    lines=15,
+                    interactive=False
+                )
+                
+                with gr.Row():
+                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
+        
+        # Event handlers
+        generate_btn.click(
+            fn=generate_text,
+            inputs=[
+                prompt_input,
+                model_dropdown,
+                max_length_slider,
+                temperature_slider,
+                top_k_slider
+            ],
+            outputs=output_text
+        )
+        
+        clear_btn.click(
+            fn=lambda: ("", ""),
+            outputs=[prompt_input, output_text]
+        )
+        
+        # Model information
+        gr.Markdown("## ℹ️ Model Information")
+        gr.Markdown("""
+        - **Model Architecture**: Kimi K2 inspired (StoryKimi)
+        - **ZeroGPU**: Dynamic GPU allocation with H200 slice (70GB VRAM)
+        - **GPU Duration**: 120 seconds maximum per generation
+        - **Deployment**: Hugging Face Spaces with automatic scaling
+        """)
+        
+        gr.Markdown("## 🚀 Features")
+        gr.Markdown("""
+        - **Top-k Sampling**: Control randomness with top-k token selection
+        - **Temperature Control**: Adjust creativity vs coherence
+        - **Variable Length**: Generate 10-128 tokens
+        - **Real-time Generation**: Powered by ZeroGPU infrastructure
+        """)
+    
+    return demo
+
+if __name__ == "__main__":
+    # Create and launch the interface
+    demo = create_interface()
+    demo.launch()

config.py

@@ -0,0 +1,151 @@
+import argparse
+from dataclasses import dataclass
+
+def get_args():
+    parser = argparse.ArgumentParser(description='SmolKimi - DeepSeek V3 Inspired Model Training')
+    
+    # Model Architecture
+    parser.add_argument('--block_size', type=int, default=128, help='Maximum sequence length')
+    parser.add_argument('--batch_size', type=int, default=256, help='Training batch size')
+    parser.add_argument('--embeddings_dims', type=int, default=384, help='Model embedding dimensions')
+    parser.add_argument('--no_of_heads', type=int, default=8, help='Number of attention heads')
+    parser.add_argument('--no_of_decoder_layers', type=int, default=6, help='Number of decoder layers')
+    parser.add_argument('--latent_dim', type=int, default=64, help='Latent dimension for attention')
+    
+    # MoE Configuration
+    parser.add_argument('--experts', type=int, default=8, help='Number of MoE experts')
+    parser.add_argument('--top_experts', type=int, default=2, help='Number of experts to route to (top-k)')
+    parser.add_argument('--use_shared_expert', action='store_true', default=True, help='Enable shared expert in MoE')
+    parser.add_argument('--noisy_topk', action='store_true', default=False, help='Use noisy top-k routing')
+    parser.add_argument('--useauxFreeLoadBalancingLoss', action='store_true', default=True, help='Use auxiliary-free load balancing loss')
+    parser.add_argument('--aux_free_bias_update_rate', type=float, default=0.001, help='Bias update rate for load balancing')
+    parser.add_argument('--loss_scale', type=float, default=0.3, help='Loss scaling factor')
+    
+    # Training Hyperparameters
+    parser.add_argument('--epochs', type=int, default=1, help='Number of training epochs')
+    parser.add_argument('--max_lr', type=float, default=6e-4, help='Maximum learning rate')
+    parser.add_argument('--weight_decay_optim', type=float, default=0.1, help='Weight decay for optimizer')
+    parser.add_argument('--beta_1', type=float, default=0.9, help='Beta1 for optimizer')
+    parser.add_argument('--beta_2', type=float, default=0.95, help='Beta2 for optimizer')
+    parser.add_argument('--eps', type=float, default=1e-8, help='Epsilon for optimizer')
+    parser.add_argument('--clip', type=float, default=1.0, help='Gradient clipping value')
+    
+    # Regularization
+    parser.add_argument('--dropout', type=float, default=0.1, help='Dropout rate')
+    parser.add_argument('--attn_dropout', type=float, default=0.1, help='Attention dropout rate')
+    
+    # System Configuration
+    parser.add_argument('--device', type=str, default='cuda', help='Device to use (cuda/cpu)')
+    parser.add_argument('--use_checkpointing', action='store_true', default=False, help='Use gradient checkpointing')
+    parser.add_argument('--use_liger', action='store_true', default=True, help='Use Liger kernels for optimization')
+    parser.add_argument('--ignore_pad_token_in_loss', action='store_true', default=True, help='Ignore padding tokens in loss calculation')
+    
+    # Data Configuration
+    parser.add_argument('--vocab_size', type=int, default=32000 + 1 , help='Vocabulary size (updated based on tokenizer)')
+    parser.add_argument('--base_freq', type=int, default=100000, help='Base frequency for positional encoding')
+    parser.add_argument('--hf_token', type=str, default=None, help='Hugging Face token for accessing gated models like Llama-2')
+    
+    # Dataset Selection
+    parser.add_argument('--dataset', type=str, default='tinystories', choices=['tinystories', 'fineweb', 'tinyshakespeare'], help='Dataset to use for training')
+    
+    # Generation Parameters
+    parser.add_argument('--generation_max_length', type=int, default=50, help='Maximum length for text generation')
+    parser.add_argument('--generation_top_k', type=int, default=50, help='Top-k value for sampling during generation')
+    parser.add_argument('--generation_temperature', type=float, default=1.0, help='Temperature for sampling during generation')
+    
+    # Logging and Checkpointing
+    parser.add_argument('--log_interval', type=int, default=100, help='Steps between logging')
+    parser.add_argument('--save_interval', type=int, default=2000, help='Steps between saving checkpoints')
+    parser.add_argument('--eval_interval', type=int, default=400, help='Steps between evaluation')
+    parser.add_argument('--eval_iters', type=int, default=400, help='Number of iterations for evaluation')
+    parser.add_argument('--warmup_iters', type=int, default=400, help='Number of warmup iterations')
+    parser.add_argument('--total_iters', type=int, default=10000, help='Total training iterations')
+    parser.add_argument('--lr_decay_iters', type=int, default=10000, help='Learning rate decay iterations')
+    parser.add_argument('--wandb_project', type=str, default='smolkimi', help='Wandb project name')
+    parser.add_argument('--wandb_run_name', type=str, default=None, help='Wandb run name')
+    
+    # Batch Size Configuration
+    parser.add_argument('--total_batch_size', type=int, default=524288, help='Total batch size for gradient accumulation')
+    parser.add_argument('--micro_batch_size', type=int, default=None, help='Micro batch size (defaults to batch_size)')
+    
+    # Distributed Training
+    parser.add_argument('--use_ddp', action='store_true', default=False, help='Use distributed data parallel')
+    
+    return parser.parse_args()
+
+@dataclass
+class ModelArgs:
+    def __init__(self, args=None):
+        if args is None:
+            args = get_args()
+        
+        # Model Architecture
+        self.block_size = args.block_size
+        self.batch_size = args.batch_size
+        self.embeddings_dims = args.embeddings_dims
+        self.no_of_heads = args.no_of_heads
+        self.no_of_decoder_layers = args.no_of_decoder_layers
+        self.latent_dim = args.latent_dim
+        
+        # MoE Configuration
+        self.experts = args.experts
+        self.top_experts = args.top_experts
+        self.use_shared_expert = args.use_shared_expert
+        self.noisy_topk = args.noisy_topk
+        self.useauxFreeLoadBalancingLoss = args.useauxFreeLoadBalancingLoss
+        self.aux_free_bias_update_rate = args.aux_free_bias_update_rate
+        self.loss_scale = args.loss_scale
+        
+        # Training Hyperparameters
+        self.epochs = args.epochs
+        self.max_lr = args.max_lr
+        self.weight_decay_optim = args.weight_decay_optim
+        self.beta_1 = args.beta_1
+        self.beta_2 = args.beta_2
+        self.eps = args.eps
+        self.clip = args.clip
+        
+        # Regularization
+        self.dropout = args.dropout
+        self.attn_dropout = args.attn_dropout
+        
+        # System Configuration
+        self.device = args.device
+        self.use_checkpointing = args.use_checkpointing
+        self.use_liger = args.use_liger
+        self.ignore_pad_token_in_loss = args.ignore_pad_token_in_loss
+        
+        # Data Configuration
+        self.vocab_size = args.vocab_size
+        self.base_freq = args.base_freq
+        self.hf_token = args.hf_token
+        self.dataset = args.dataset
+        
+        # Generation Parameters
+        self.generation_max_length = args.generation_max_length
+        self.generation_top_k = args.generation_top_k
+        self.generation_temperature = args.generation_temperature
+        
+        # Logging and Checkpointing
+        self.log_interval = args.log_interval
+        self.save_interval = args.save_interval
+        self.eval_interval = args.eval_interval
+        self.eval_iters = args.eval_iters
+        self.warmup_iters = args.warmup_iters
+        self.total_iters = args.total_iters
+        self.lr_decay_iters = args.lr_decay_iters
+        self.wandb_project = args.wandb_project
+        self.wandb_run_name = args.wandb_run_name
+        
+        # Batch Size Configuration
+        self.total_batch_size = args.total_batch_size
+        self.micro_batch_size = args.micro_batch_size if args.micro_batch_size else args.batch_size
+        self.gradient_accumulation_steps = self.total_batch_size // (self.micro_batch_size * self.block_size)
+        
+        # Calculated parameters
+        self.min_lr = 0.1 * self.max_lr
+        self.save_checkpoint_iter = self.save_interval
+        self.eval_check = self.eval_interval
+        
+        # Distributed Training
+        self.use_ddp = args.use_ddp

inference.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn.functional as F
+from config import ModelArgs
+from model import DeepSeekV3
+from tokenizer import Tokenizer
+
+def topk_sampling(model, prompt, device, max_length=50, top_k=50, temperature=1.0, tokenizer=None, hf_token=None):
+    if tokenizer is None:
+        # Use default tokenizer if none provided
+        tokenizer_instance = Tokenizer(hf_token=hf_token)
+        tokenizer = tokenizer_instance.ready_tokenizer()
+        
+    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
+    generated_tokens = []
+    
+    if(len(input_ids[0]) < max_length):
+        max_length -= len(input_ids[0]) # If the input is longer than max_length, set max_length to the length of the input
+    else:
+        max_length = len(input_ids[0]) - max_length  
+    for _ in range(max_length):
+        with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+            # Pass inference=True to use the inference path in the model
+            outputs = model(input_ids, inference=True)
+            logits = outputs[:, -1, :]
+            logits = logits / temperature
+            probs = F.softmax(logits, dim=-1)
+            
+            # Top-k filtering
+            top_k_probs, top_k_indices = torch.topk(probs, top_k, dim=-1)
+            
+            # Sample from top-k
+            next_token = torch.multinomial(top_k_probs, num_samples=1)
+            
+            xcol = torch.gather(top_k_indices, -1, next_token)
+            input_ids = torch.cat([input_ids, xcol], dim=1) #1 because is it the dimension of the sequence
+            
+            if hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id and xcol.item() == tokenizer.eos_token_id:
+                break
+            
+            
+    return tokenizer.decode(input_ids[0])
+
+
+def save_text(file_path, step, text):
+    with open(file_path, 'w') as f:
+        f.write(f"Step {step}: {text}\n")

model.py

@@ -0,0 +1,589 @@
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from torch.nn import RMSNorm
+from config import ModelArgs
+from tokenizer import Tokenizer
+
+# Initialize tokenizer globally as None - will be set later
+tokenizer = None
+model_args = ModelArgs()
+
+
+def initialize_tokenizer(hf_token=None):
+    """Initialize the global tokenizer with the provided HF token"""
+    global tokenizer
+    if tokenizer is None:
+        tokenizer_instance = Tokenizer(hf_token=hf_token)
+        tokenizer = tokenizer_instance.ready_tokenizer()
+    return tokenizer
+
+class Normalization(nn.Module):
+    def __init__(
+        self,
+        embeddings_dims: int = model_args.embeddings_dims
+    ):  
+        super().__init__()
+        self.rmsnorm_layer = RMSNorm(embeddings_dims)
+        
+        
+    def forward(self, x):
+        
+        x = self.rmsnorm_layer(x)
+        return x
+        
+
+
+class Swish(nn.Module):
+    def __init__(
+        self,
+        block_size: int = model_args.block_size,
+        embeddings_dims: int = model_args.embeddings_dims,
+        device = model_args.device
+    ):
+        super().__init__()
+
+        self.sig = torch.nn.Sigmoid()
+
+
+    def forward(self, x):
+        swish = x * self.sig(x)
+
+        return swish
+
+
+
+class SWiGLUExpertMoE(nn.Module):
+    def __init__(
+        self,
+        block_size: int = model_args.block_size,
+        embeddings_dims: int = model_args.embeddings_dims,
+        device = model_args.device
+    ):
+        super().__init__()
+
+        self.hidden_dims = (embeddings_dims * 2)
+        self.swish = Swish(block_size=block_size, embeddings_dims=embeddings_dims, device=device)
+        self.linear_layer1 = nn.Linear(in_features=embeddings_dims, out_features=self.hidden_dims,  bias=False, device = device)
+        self.linear_layer2 = nn.Linear(in_features=embeddings_dims, out_features=self.hidden_dims,  bias=False, device = device)
+        self.linear_layer3 = nn.Linear(in_features=self.hidden_dims, out_features=embeddings_dims,  bias=False, device = device)
+
+
+
+
+    def forward(self, x):
+        swish_res = self.swish(self.linear_layer1(x))
+        x_V = self.linear_layer2(x)
+        res = torch.mul(swish_res, x_V)
+        out = self.linear_layer3(res)
+        return out
+
+
+
+class MoeLayer(nn.Module):
+    def __init__(
+        self,
+        dropout = model_args.dropout,
+        embeddings_size = model_args.embeddings_dims,
+        device = model_args.device,
+        # inner_dimensional_states: int = 3072
+    ):
+        super().__init__()
+
+        self.heads = nn.ModuleList([SWiGLUExpertMoE() for _ in range(model_args.experts)])
+        self.gate = nn.Linear(in_features=embeddings_size, out_features=model_args.experts, device=device, bias=False)
+        
+        # Only create shared expert if enabled
+        if model_args.use_shared_expert:
+            self.shared_expert = SWiGLUExpertMoE()
+        else:
+            self.shared_expert = None
+            
+        if(model_args.noisy_topk is True and model_args.use_checkpointing == False):
+            self.noise = nn.Linear(in_features=embeddings_size, out_features=model_args.experts, device=device, bias=False)
+            self.noisy_router = None
+        # self.outputs = torch.zeros((batch_size,block_size, embeddings_size), device=device) #batch size needs to be defined because we are accessing it explicitly
+        self.device = device
+        # self.shared_expert_out = torch.zeros((model_args.batch_size, model_args.embeddings_dims), device=device)
+        # self.b = torch.zeros((model_args.batch_size, model_args.block_size, model_args.experts), device=device)
+
+        if model_args.useauxFreeLoadBalancingLoss:
+            self.register_buffer('routing_bias', torch.zeros(model_args.experts, device=self.device))
+            # self.routing_bias = torch.zeros(model_args.experts, device=self.device)
+            self.bias_update_speed = model_args.aux_free_bias_update_rate
+        
+        
+    def forward(self, x):
+        # mlp_weights_init = self.mlp.apply(weights_init)
+        self.gate_out = self.gate(x) #[bz, seq, num_experts]
+
+        
+        if(model_args.noisy_topk == True and model_args.use_checkpointing == False):
+            noise = self.noise(x)
+            gaussian_noise = torch.normal(0, 1, size=self.gate_out.shape, device=self.device)
+            self.noisy_router = F.softplus(noise) * gaussian_noise
+            self.gate_out += self.noisy_router
+            
+        
+
+        shared_output = 0
+        out = 0
+
+        
+
+        if model_args.useauxFreeLoadBalancingLoss:
+            
+           self.gate_out += self.routing_bias
+           
+        
+                
+        
+        # Adjust top_k based on whether shared expert is used
+        top_k = model_args.top_experts
+        top_k_values, top_k_indices = torch.topk(self.gate_out, k=top_k) #[bs, seq len, top k]
+        # topkmask = torch.ones_like(top_k_values, device=self.device)  # [bs, seq len, experts]
+        # indices = torch.arange(top_k_values.size(0), device=self.device).unsqueeze(1).unsqueeze(2)  # [bs, 1, 1]
+        # topkvaluesMasked = top_k_values.masked_fill(indices != top_k_indices, float('-inf'))  # Mask out negative values
+        masked = torch.full_like(self.gate_out, float('-1e20'), device=self.device) 
+        masked_values = masked.scatter_(-1, top_k_indices, top_k_values)
+        probs = torch.nn.functional.softmax(masked_values, dim=-1) #[bs, seq len, top k]
+        
+        out = torch.zeros_like(x)
+        if model_args.use_shared_expert and self.shared_expert is not None:
+            shared_output += self.shared_expert(x)
+
+        flat_x = x.view(-1, x.size(-1))  # Flatten the input for easier processing
+
+        for i in range(model_args.experts): # Iterate through each expert index (0 to num_experts-1)
+            # Determine which tokens routed to this expert 'i'
+            # top_k_indices is [bs, seq_len, self.top_k]
+            # We want a mask of shape [bs, seq_len] where True if expert 'i' is in the top_k for that token
+            expert_i_is_chosen_mask = (top_k_indices == i).any(dim=-1) # Check along the top_k dimension
+            # expert_i_is_chosen_mask has shape [bs, seq_len]
+
+            if not expert_i_is_chosen_mask.any(): # If expert 'i' was not chosen by any token
+                continue
+
+            # Flatten the mask to apply to flat_x
+            flat_expert_i_is_chosen_mask = expert_i_is_chosen_mask.reshape(-1) # Shape: [bs * seq_len]
+
+            # Select input tokens for this expert
+            selected_input_tokens = flat_x[flat_expert_i_is_chosen_mask] # Shape: [num_active_for_expert_i, embed_dim]
+
+            if selected_input_tokens.numel() == 0: # Should be caught by .any() above, but good check
+                continue
+
+            # Process through the expert
+            expert_output_for_selected = self.heads[i](selected_input_tokens)
+
+            # Get the routing probabilities for these chosen tokens specifically for expert 'i'
+            # routing_probs is [bs, seq_len, num_experts]
+            # expert_i_probs_original_shape = routing_probs[:, :, i] # Probabilities for expert 'i', shape [bs, seq_len]
+            # flat_expert_i_probs = expert_i_probs_original_shape.reshape(-1) # Shape [bs * seq_len]
+            # active_token_weights = flat_expert_i_probs[flat_expert_i_is_chosen_mask] # Shape: [num_active_for_expert_i]
+
+            # Alternative way to get weights directly using the mask on routing_probs for expert i:
+            # Get the [bs, seq_len] slice of probabilities for the current expert 'i'
+            probs_for_expert_i = probs[:, :, i] # Shape: [bs, seq_len]
+            # Now use the expert_i_is_chosen_mask (which is also [bs, seq_len]) to select the relevant weights
+            active_token_weights = probs_for_expert_i[expert_i_is_chosen_mask] # Shape: [num_active_for_expert_i]
+
+
+            weighted_expert_output = expert_output_for_selected * active_token_weights.unsqueeze(-1)
+
+            # Add this expert's contribution
+            temp_contribution_for_expert_i = torch.zeros_like(x) # Initialize with zeros
+            temp_contribution_for_expert_i.masked_scatter_(
+                expert_i_is_chosen_mask.unsqueeze(-1).expand_as(x), # Use the original 2D mask, expanded
+                weighted_expert_output
+            )
+            out = out + temp_contribution_for_expert_i
+            
+            
+        # for expert_idx in range(model_args.experts):
+        #     # Create mask for current expert across all top_k positions
+        #     expert_mask = (top_k_indices == expert_idx)
+
+        #     # Sum probabilities for current expert
+        #     expert_weights = (probs * expert_mask).sum(dim=-1)  # [batch, seq_len]
+
+        #     # Get inputs where expert is used
+        #     selected = expert_weights > 0
+        #     if not selected.any():
+        #         continue
+        #     # print(expert_weights.shape)
+        #     # print(x[selected].shape)
+
+        #     # Process all selected inputs through expert
+        #     expert_out = self.heads[expert_idx](x[selected])
+            
+            
+                
+        #     # Weight and accumulate outputs
+        #     out[selected] += expert_out * expert_weights[selected].unsqueeze(-1)
+
+        out = out + shared_output  # Add shared expert output if enabled
+        
+        if model_args.useauxFreeLoadBalancingLoss and self.training:
+            
+            with torch.no_grad():  
+                ci = probs.sum(dim=(0,1))  # Su  of tokens for each expert
+                ci_avg = ci.mean()
+                
+                
+                error_i = ci_avg - ci
+                
+                self.update = self.bias_update_speed * torch.sign(error_i)  # Update routing bias
+                self.routing_bias.add_(self.update)
+                # self.routing_bias = self.routing_bias + self.update
+
+        return out
+    
+    
+# import numpy as np
+class SinusoidalPositionalEmbeddings(nn.Module):
+    def __init__(
+        self,
+        device,
+        embeddings_dims: int = model_args.embeddings_dims,
+        block_size: int = model_args.block_size,
+        batch_size: int = model_args.batch_size,
+    ):
+        super().__init__()
+
+        self.embeddings_dims = embeddings_dims
+        self.block_size = block_size
+        self.batch_size = batch_size
+        self.device = device
+        
+        # Create positional encoding matrix
+        pe = torch.zeros(block_size, embeddings_dims)
+        position = torch.arange(0, block_size, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, embeddings_dims, 2).float() * (-math.log(10000.0) / embeddings_dims))
+        
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        
+        # Register as buffer so it's not a parameter but moves with the model
+        self.register_buffer('pe', pe.unsqueeze(0))  # Shape: [1, block_size, embeddings_dims]
+
+    def forward(self, x):
+        # x shape: [batch_size, seq_len, embeddings_dims]
+        batch_size, seq_len, _ = x.shape
+        
+        # Add positional embeddings
+        # pe[:, :seq_len] ensures we only use the positional embeddings up to the sequence length
+        pos_emb = self.pe[:, :seq_len].to(x.device)
+        return pos_emb
+
+
+
+class LatentAttention(nn.Module):
+    def __init__(
+        self,
+        attn_dropout = model_args.attn_dropout,
+        embeddings_dims = model_args.embeddings_dims,
+        no_of_heads = model_args.no_of_heads,
+        device = model_args.device
+    ):
+        super().__init__()
+        self.head_size = embeddings_dims // no_of_heads
+        self.no_of_heads = no_of_heads
+        # if(model_args.use_flash_attention==False):
+        self.latent_dim = model_args.latent_dim
+        self.W_k = nn.Linear(in_features=self.latent_dim, out_features=self.head_size, device=device, bias=False)
+        self.W_v = nn.Linear(in_features=self.latent_dim, out_features=self.head_size, device=device, bias=False)
+        self.W_dkv = nn.Linear(in_features=model_args.embeddings_dims, out_features=self.latent_dim, device=device, bias=False) # 3 for query, key and value
+        self.query = nn.Linear(in_features=embeddings_dims, out_features=self.head_size, device=model_args.device, bias=False)
+        # self.keys = nn.Linear(in_features=embeddings_dims, out_features=self.head_size,device=model_args.device, bias=False)
+        # self.values = nn.Linear(in_features=embeddings_dims, out_features=self.head_size, device=model_args.device,bias=False)
+    # self.dropout = nn.Dropout(p = attn_dropout)
+        
+
+        self.dropout = nn.Dropout(p = attn_dropout)
+        self.device = device
+
+        # Use sinusoidal positional embeddings instead of rotary
+        self.pos_embeddings = SinusoidalPositionalEmbeddings(embeddings_dims=self.head_size, device=device)
+        # self.register_buffer('absorbed_q', None)
+        # self.absorbed_q = None
+        
+    def forward(self, x, kv_cache=None, mask=None):
+        batch_size, block_size, embd_dims = x.shape
+
+        # k = self.keys(x)
+        # q = self.query(x)
+        # v = self.values(x)
+        
+        self.latent_matrix = self.W_dkv(x)
+
+        # print("q shape: ", q.shape)
+        
+        # print("Shape of latent mat: ", self.query.weight.shape)
+        # print("Shape of compressed_k: ", self.W_k.weight.shape)
+        
+        # if(self.absorbed_q is None):
+        self.absorbed_q = torch.matmul(self.query.weight.T , self.W_k.weight)
+        
+        
+        # weights = q @ torch.transpose(k, dim0=-2, dim1=-1) * (k.shape[-1] ** -0.5)
+
+        # if kv_cache is None:
+        if kv_cache is None:
+            kv_cache = self.latent_matrix
+        else:
+            # print(kv_cache)
+            # print("Shape of latent matrix: ", self.latent_matrix.shape)
+            # print("Shape of kv_cache: ", kv_cache.shape)
+            kv_cache = torch.cat([kv_cache, self.latent_matrix], dim=1)
+
+        self.compressed_k = self.W_k(kv_cache)
+        self.compressed_v = self.W_v(kv_cache)
+        
+        q_res = torch.matmul(x , self.absorbed_q)
+        weights =  q_res @ torch.transpose(kv_cache, dim0=-2, dim1=-1) * (self.head_size ** -0.5)  # [batch_size, block_size, block_size]
+        # print("Shape of weights: ", weights.shape)
+        # print("Shape of kv_cache: ", kv_cache.shape)
+        if(mask is not None):
+            weights = weights.masked_fill(mask == 0, float('-1e20')) #Masking the attention weights
+            
+        masked_table = torch.tril(torch.ones(q_res.shape[1], kv_cache.shape[1], device=model_args.device))
+
+        masked_values = weights.masked_fill(masked_table[: q_res.shape[1], : kv_cache.shape[1]] == 0, float('-1e20'))
+        weights_normalized = nn.functional.softmax(masked_values, dim=-1) #Normalize along the embeddings dimension for all the tokens
+        weights_normalized = self.dropout(weights_normalized)
+        
+        # print("Shape of weights_normalized: ", weights_normalized.shape)
+        # Apply positional embeddings to the output
+        
+        
+        
+        
+        # print("Shape of compressed_v: ", self.compressed_v.shape)
+        out = weights_normalized @ self.compressed_v
+        
+        # out = self.pos_embeddings(out)
+        return out, kv_cache
+
+# MHA
+
+
+class MHLA(nn.Module):
+    def __init__(
+        self,
+        device,
+        attn_dropout = model_args.attn_dropout,
+        embeddings_dims = model_args.embeddings_dims,
+        no_of_heads = model_args.no_of_heads,
+    ):
+        super().__init__()
+        self.heads = nn.ModuleList([LatentAttention(attn_dropout=attn_dropout, embeddings_dims=embeddings_dims, no_of_heads=no_of_heads) for _ in range(no_of_heads)])
+        self.dropout = nn.Dropout(p = attn_dropout)
+        self.linear = nn.Linear(in_features=embeddings_dims, out_features=embeddings_dims, device=device, bias=False) # 12 (no of heads) * (batch_size) 64 = 768 -> gives out the text embeddings
+        
+    def forward(self, x, kv_cache=None, mask=None):
+        # concat = torch.cat([head(x, kv_cache=kv_cache, mask=mask) for head in self.heads], dim=-1)
+        res = []
+        for head in self.heads:
+            head_out, kv_cache = head(x, kv_cache=kv_cache, mask=mask)
+            res.append(head_out)
+        concat = torch.cat(res, dim=-1)  # Concatenate along the last dimension
+        linear_layer = self.linear(concat)
+        out = self.dropout(linear_layer)
+        return out, kv_cache
+
+class FFN(nn.Module):
+    def __init__(self,
+                  device,
+                  embeddings_dims: int = model_args.embeddings_dims,
+                  block_size: int = model_args.block_size,
+                  vocab_size: int = model_args.vocab_size,
+                   dropout = model_args.dropout
+
+                 ):
+        super().__init__()
+
+        self.linear_layer = nn.Linear(in_features=embeddings_dims, out_features=embeddings_dims,  dtype=torch.float32,  device = device)
+        self.linear_layer2 = nn.Linear(in_features=embeddings_dims, out_features=embeddings_dims,  dtype=torch.float32, device = device)
+
+        self.dropout = nn.Dropout(p = dropout)  # Uncommenting the dropout line
+    def forward(self, x):
+
+        x = self.linear_layer(x)
+        x = F.gelu(x)
+        x = self.linear_layer2(x)
+        x = F.gelu(x)
+        # x = self.dropout(x)  # Uncommenting the dropout line
+        return x
+
+
+
+
+
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self,
+                device,
+                attn_dropout: float = model_args.attn_dropout,
+                no_of_heads: int = model_args.no_of_heads,
+                embeddings_dims: int = model_args.embeddings_dims,
+                dropout = model_args.dropout,
+                block_size: int = model_args.block_size,
+                vocab_size: int = model_args.vocab_size,
+
+                 ) :
+        super().__init__()
+
+        # self.base_freq = model_args.base_freq
+        # self.feedforward_network = FFN(embeddings_dims=embeddings_dims, block_size=block_size, vocab_size=vocab_size,  device = device)
+        self.mha = MHLA(attn_dropout=attn_dropout, embeddings_dims=embeddings_dims, no_of_heads=no_of_heads, device=device)
+        self.layer_norm1 = Normalization(embeddings_dims=embeddings_dims)
+        self.layer_norm2 = Normalization(embeddings_dims=embeddings_dims)
+        # self.layer_norm3 = Normalization(embeddings_dims=embeddings_dims)
+        self.dropout = nn.Dropout(p = dropout)
+
+        self.moe_block = MoeLayer(dropout=dropout, embeddings_size=embeddings_dims)
+
+    def forward(self, x, kv_cache=None, ffn=None, mask=None):
+
+        out, kv_cache = self.mha(self.layer_norm1(x), kv_cache=kv_cache, mask=mask)  #Very important step -> Layer Norm on input and then passes it to the subsequent blocks
+        x = x + out  # Fixed: removed in-place operation
+        x = x + self.moe_block(self.layer_norm2(x)) #Very important step
+
+        return x, kv_cache
+
+
+class Block(nn.Module):
+    def __init__(self,
+                  device,
+                  embeddings_dims: int = model_args.embeddings_dims,
+                  no_of_decoder_layers: int = model_args.no_of_decoder_layers,
+                  block_size: int = model_args.block_size,
+                  vocab_size: int = model_args.vocab_size,
+                  dropout = model_args.dropout
+
+                 ) :
+        super().__init__()
+        self.base_freq = model_args.base_freq
+        # self.embeddings = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embeddings_dims,  dtype=torch.float32,  device = device)
+        self.decoder = nn.ModuleList(DecoderLayer(embeddings_dims=embeddings_dims, block_size=block_size, vocab_size=vocab_size, dropout=dropout,  device = device) for _ in range(no_of_decoder_layers))
+        # self.linear_layer = nn.Linear(in_features=embeddings_dims, out_features=vocab_size,  dtype=torch.float32,  device = device)
+        self.dropout = nn.Dropout(p = dropout)
+        self.norm = Normalization(embeddings_dims)
+        
+        #weight tying
+        # self.embeddings.weight = self.linear_layer.weight
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+            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, x, mask=None, actual_labels = None, inference=False):
+        index = 0
+        no_of_layers = 0
+        # x = self.embeddings(x)
+        # # x = self.dropout(x)
+        # if(mask is not None):
+        kv_cache = None
+        #     x = x * mask
+        #     # mask = mask.unsqueeze(-1)
+        # x = self.decoder(x)
+        for layer in self.decoder:
+            # if no_of_layers % 2 == 0:
+            #     if no_of_layers % 4 == 0:
+            #         # print("x shape: ", x.shape)
+            #         x = layer(x, rope=False, ffn=True, mask=mask)
+            #     x = layer(x, rope=True, ffn=True, mask=mask)
+                
+            #     # print("x shape: ", x.shape)
+            # else:
+            #     # print("x shape local: ", x.shape)
+            #     if no_of_layers % 4 == 0:
+            #         # print("x shape: ", x.shape)
+            #         x = layer(x, rope=False, ffn=False, mask=mask)
+            x, kv_cache = layer(x, kv_cache=kv_cache, ffn=None, mask=mask)
+                # print("x shape local: ", x.shape)
+            # no_of_layers += 1
+        # print(x.shape)
+        x = self.dropout(x)
+        x = 2 * ((model_args.no_of_decoder_layers) ** -0.5) * x
+        x = self.norm(x)
+        
+        # if(inference):
+        #     out = self.linear_layer(x)
+        #     return out
+        # if(model_args.use_liger):  
+        #     # print("yo")
+        #     y = x.contiguous().view(-1, model_args.embeddings_dims)
+        #     if(actual_labels is not None):
+        #         labels = actual_labels.contiguous().view(-1)
+                
+        #         # Pass linear layer weights FIRST as required [2][5]
+        #         # ignore_index is already set during initialization
+        #         loss = self.le_loss(self.linear_layer.weight, y, labels)
+        #         return loss
+        # else:
+        #     # print("Hi")
+        #     out = self.linear_layer(x)
+        #     return out
+
+        return x
+
+
+
+class DeepSeekV3(nn.Module):
+    def __init__(self,
+                 device,
+                 embeddings_dims: int = model_args.embeddings_dims,
+                 block_size: int = model_args.block_size,
+                 vocab_size: int = model_args.vocab_size,
+                 dropout = model_args.dropout
+                ):
+        super().__init__()
+        self.decoder = Block(device=device, embeddings_dims=embeddings_dims, no_of_decoder_layers=model_args.no_of_decoder_layers, block_size=block_size, vocab_size=vocab_size, dropout=dropout)
+        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embeddings_dims, dtype=torch.float32, device=device)
+        self.pos_embeddings = SinusoidalPositionalEmbeddings(embeddings_dims=embeddings_dims, device=device)
+        self.linear_layer = nn.Linear(in_features=embeddings_dims, out_features=vocab_size, dtype=torch.float32, device=device, bias=False)
+        # Weight tying - tie embedding and output projection weights
+        self.embedding.weight = self.linear_layer.weight
+        
+        # Initialize the LigerFusedLinearCrossEntropyLoss for optimized training
+        if model_args.use_liger:
+            from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
+            # Initialize with ignore_index for padding tokens if enabled
+            if model_args.ignore_pad_token_in_loss:
+                self.le_loss = LigerFusedLinearCrossEntropyLoss(
+                    ignore_index=tokenizer.pad_token_id
+                )
+            else:
+                self.le_loss = LigerFusedLinearCrossEntropyLoss()
+
+    def forward(self, x, inference=False, mask=None):
+        if(mask is not None):
+            x = x * mask
+            
+        x = self.embedding(x)
+        x = x + self.pos_embeddings(x)  # Add positional embeddings
+        B, T, C = x.shape
+            
+        if inference:
+            # For inference, we only need the last token prediction
+            decoder_out = self.decoder(x, mask=mask)
+            logits = self.linear_layer(decoder_out)
+            return logits
+        else:
+            decoder_out = self.decoder(x, mask=mask)
+            logits = self.linear_layer(decoder_out)
+            return logits

requirements.txt

@@ -0,0 +1,9 @@
+spaces
+torch>=2.1.2
+transformers>=4.36.0
+datasets
+tqdm
+huggingface_hub
+gradio
+numpy
+safetensors

tokenizer.py

@@ -0,0 +1,18 @@
+from transformers import AutoTokenizer
+
+class Tokenizer:
+    
+    def __init__(self, hf_token=None) -> None:
+        # Try to get token from environment if not provided
+        
+        if hf_token:
+            print(f"[INFO] Using HF token for model access")
+        else:
+            print("[INFO] No HF token provided - using public models only")
+            
+        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", token=hf_token)
+        self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+
+    def ready_tokenizer(self):
+        
+        return self.tokenizer