Docs: Updated README

README.md

@@ -11,4 +11,70 @@ license: mit
 short_description: Text generation using smollmv2-135M model
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SmoLLMv2: A Small but Efficient Language Model
+
+[Training Repo Link](https://github.com/Shilpaj1994/SmoLLMv2)
+[Gradio App Link](https://huggingface.co/spaces/Shilpaj/SmoLLMv2)
+
+
+SmoLLMv2 is a 135M parameter language model designed for efficient text generation. It incorporates several modern architectural improvements while maintaining a small footprint.
+
+
+
+## Features
+
+- **Efficient Architecture**:
+  - 30 transformer layers
+  - 9 attention heads
+  - 576 embedding dimension
+  - Memory-efficient attention with reduced KV dimensions
+  - Rotary Position Embeddings (RoPE)
+  - SwiGLU activation function
+
+- **Training Optimizations**:
+  - Mixed precision training (16-bit)
+  - Gradient accumulation
+  - OneCycleLR scheduler
+  - Streaming dataset support
+  - Automatic model compilation (with PyTorch 2.0+)
+
+
+
+## Model Architecture
+
+SmoLLMv2 incorporates several efficiency improvements:
+
+1. **Reduced KV Dimensions**: Uses 189-dimensional key/value projections (instead of full 576) to save memory and computation.
+2. **RoPE Attention**: Implements Rotary Position Embeddings for better handling of sequential information.
+3. **SwiGLU Activation**: Uses the SwiGLU activation function in the MLP layers for better performance.
+4. **Weight Sharing**: Shares weights between input embeddings and output projection.
+
+
+
+## Configuration
+
+The model's behavior can be customized through various configuration classes in `config.py`:
+
+- `SmollmConfig`: Core model architecture and training parameters
+- `RoPEConfig`: Rotary Position Embedding settings
+- `OptimizerConfig`: Optimization and learning rate settings
+- `DataConfig`: Dataset and tokenizer configuration
+- `TrainerConfig`: Training infrastructure settings
+
+
+
+## Dataset
+
+The model is trained on the Cosmopedia dataset, which is streamed during training to handle large-scale data efficiently.
+
+
+
+## Requirements
+
+See `requirements.txt` for full dependencies. Key requirements:
+
+- PyTorch ≥ 2.0.0
+- Transformers ≥ 4.30.0
+- Lightning ≥ 2.0.0
+- Gradio ≥ 5.13.1
+

app.py

@@ -95,6 +95,11 @@ model, tokenizer, device = load_model()
 def generate_text(prompt, num_tokens, temperature=0.8, top_p=0.9):
     """
     Generate text using the SmollmV2 model.
+    :param prompt: The initial text prompt to start the generation from.
+    :param num_tokens: The number of tokens to generate.
+    :param temperature: The temperature parameter for controlling randomness.
+    :param top_p: The top-p parameter for nucleus sampling
+    :return: The generated text.
     """
     try:
         # Ensure num_tokens doesn't exceed model's block size
@@ -148,13 +153,13 @@ demo = gr.Interface(
     fn=generate_text,
     inputs=[
         gr.Textbox(label="Enter your prompt", value="Once upon a time"),
-        gr.Slider(minimum=1, maximum=SmollmConfig.block_size, value=100, step=1, label="Number of tokens to generate"),
+        gr.Slider(minimum=1, maximum=SmollmConfig.block_size//2, value=100, step=1, label="Number of tokens to generate"),
         gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature (higher = more random)"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p (nucleus sampling)")
     ],
     outputs=gr.Textbox(label="Generated Text"),
-    title="SmollmV2 Text Generator",
-    description="Generate text using the SmollmV2 model",
+    title="SmoLLMv2 Text Generator",
+    description="Generate text using the SmoLLMv2-135M model",
     allow_flagging="never",
     cache_examples=True
 )