Update README.md

README.md

@@ -15,9 +15,10 @@ tags:
 - 5k-context
 - Supra
 - Supra-50M
+- chatml
 ---
 
-<h1 align="center">Supra1.5-50M Base</h1>
+<h1 align="center">Supra1.5-50M Base (ChatML)</h1>
 
 <p align="center">
   Continued Pretraining • 50M Parameters • 5K Context
@@ -25,33 +26,39 @@ tags:
 
 ![Supra-1.5-Base-EXP](https://cdn-uploads.huggingface.co/production/uploads/68a5d0966d33a07f8aad2e51/Lh7KcQs60Ht9iray8WFbp.png)
 
-Supra-1.5-50M-Base-exp is a continued-pretrained 50M parameter Llama-style base
-model derived from `SupraLabs/Supra-50M-Base`. The target update expands the
-usable context window from 1,024 tokens to 5,120 tokens using RoPE scaling and
-full-weight continued pretraining.
+Supra-1.5-50M-Base-exp is a continued-pretrained 50M parameter Llama-style base model derived from `SupraLabs/Supra-50M-Base`. The context window has been expanded from 1,024 to 5,120 tokens using RoPE scaling and full-weight continued pretraining.
 
+This specific repository has been modified to natively support ChatML. The tokenizer and embedding layers have been updated to include `<|im_start|>` and `<|im_end|>` as single special tokens.
 
-
-## Architecture
-
-The model keeps the original Supra-50M architecture and tokenizer:
+## Architecture & Updates
 
 | Specification | Value |
 |--------------|--------|
 | Architecture | `LlamaForCausalLM` |
 | Parameters | ~50M |
-| Vocabulary Size | 32,000 |
+| Vocabulary Size | 32,002 (Expanded for ChatML) |
 | Hidden Size | 512 |
 | Layers | 12 |
 | Attention Heads | 8 |
 | KV Heads | 4 |
 | Context Length | 5,120 tokens |
-| Tokenizer | Original Supra byte-level BPE tokenizer |
+| Tokenizer | BPE tokenizer (ChatML Jinja template applied) |
+| Native EOS Token | `<|im_end|>` |
+
+## Vocabulary Expansion & Initialization
+
+The model's 3B token CPT corpus included raw ChatML text. As a result, the base model originally learned the tags `<|im_start|>` and `<|im_end|>` as sequences of 7 separate subwords. 
+
+To convert these tags into single tokens without destroying the learned pre-trained representations, we used subword-mean initialization:
+1. Extracted the attention embeddings for the original 7-token sequences.
+2. Averaged the weights for each sequence.
+3. Assigned these mean weights to the new single-token IDs.
+
+This ensures stable loss at the start of fine-tuning and allows inference frameworks to natively stop generation when `<|im_end|>` is predicted.
 
 ## Continued Pretraining Objective
 
-This is CPT, not instruction fine-tuning. Training uses packed raw text with
-standard causal language-modeling loss:
+This is a base model, not an instruct model. Training used packed raw text with standard causal language-modeling loss:
 
 - `labels = input_ids`
 - all non-pad tokens are trained
@@ -61,7 +68,7 @@ standard causal language-modeling loss:
 
 ## Data Mix
 
-The current local training mix prepared for this run is:
+The training mix prepared for the CPT run was:
 
 - 3,000,000,062 CPT tokens
   - 30% Tool Calling
@@ -69,6 +76,14 @@ The current local training mix prepared for this run is:
   - 25% Factual Text (articles, essays, blogs)
   - 15% Math & Logic Questions
 
+## Fine-Tuning Guide
+
+This model is intended for Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL), including reasoning-style (R1) alignments. 
+
+You can use standard Hugging Face tools and call `tokenizer.apply_chat_template()` directly on your datasets.
+
+**LoRA Configuration Note:** 
+Because the vocabulary size was expanded from 32,000 to 32,002, you must train the embedding layers during fine-tuning so the model can update the new ChatML tokens. Add the following to your PEFT/LoRA configuration:
 
-### Intended Use
-Supervised Fine-Tuning (SFT) and Reinforcement Learning
+```python
+modules_to_save=["embed_tokens", "lm_head"]