Updated README.md

README.md

@@ -1,3 +1,139 @@
 ---
+library_name: transformers
 license: apache-2.0
+datasets:
+- HuggingFaceTB/smoltalk
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- causal-lm
+- fine-tuned
+- instruct-model
+- custom-architecture
+- pytorch
+- tiktoken
+- chatml
 ---
+
+<p align="center">
+  <img src="./axiom_logo.png" width="220">
+</p>
+
+# Axiom-Dense-380M-Instruct
+
+Axiom-Dense-380M-Instruct is a fine-tuned, instruction-following decoder-only causal language model. It was trained by performing Supervised Fine-Tuning (SFT) on the base model [Axiom-Dense-380M-Base](https://huggingface.co/user-anto/Axiom-Dense-380M-Base) using instruction-response conversational data.
+
+# Quickstart
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_name = "user-anto/Axiom-Dense-380M-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu")
+
+prompt = "<|im_start|>user\nWrite a short email to my team about meeting tomorrow.<|im_end|>\n<|im_start|>assistant\n"
+inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+
+with torch.no_grad():
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=128,
+        temperature=0.2,
+        top_p=0.85,
+        repetition_penalty=1.15,
+        no_repeat_ngram_size=3,
+    )
+
+print(tokenizer.decode(outputs[0]))
+```
+
+## Model Summary
+
+- Model type: decoder-only Transformer (causal LM)
+- Parameter count: 385,849,344
+- Context length: 1,024 tokens
+- Vocabulary: 100,277 (`tiktoken` `cl100k_base` with ChatML special tokens patched)
+- Training objective: Autoregressive supervised fine-tuning (SFT) using target masking (only computing loss on the assistant's responses)
+- Prompt format: ChatML (`<|im_start|>`, `<|im_end|>`)
+
+## Architecture
+
+This model preserves the same dense Transformer stack as the base model, but utilizes added special tokens to delimit speaker turns during inference.
+
+- Hidden size: 1024
+- Layers: 24
+- Attention heads: 16
+- KV heads: 8 (GQA)
+- FFN multiplier: 2.6667 (rounded to 2816 intermediate dimension)
+- Normalization: RMSNorm
+- Positional encoding: RoPE (`theta=10000`)
+- Activation: SwiGLU
+- Special tokens: `<|im_start|>` (100264) and `<|im_end|>` (100265) for ChatML boundaries
+
+## Training Data
+
+- Source dataset: `HuggingFaceTB/smoltalk`
+- Local dataset path during training: `data/smol-smoltalk`
+- SFT targets: Computes loss only on assistant response tokens, masking out prompt and user tokens.
+- Total training tokens: 204,802,175 (~0.205B tokens)
+- Validation tokens: 197,825 tokens
+
+## SFT Training Setup
+
+- Effective tokens per optimizer step: 319,488 (`batch_size=1`, `seq_len=1024`, `grad_accum=312`)
+- Total optimizer steps: 641
+- Optimizer: AdamW8bit (with bitsandbytes)
+- LR schedule: warmup, constant phase, cosine decay
+- Warmup steps: 51 steps (8% of training)
+- Cosine decay phase: 102 steps (16% of training, starting at step 539)
+- LR max/min: 3e-4 / 3e-5 (initial learning rate starts at 1.5e-4 during warmup)
+- Weight decay: 0.1
+- Precision: bfloat16
+- Gradient checkpointing: enabled
+
+## Evaluation Snapshot
+
+- Pretraining base perplexity: 18.1233
+- Best observed SFT eval loss: 1.2641 at step 630
+- Best observed SFT eval perplexity: 3.5398 at step 630
+- Final SFT step (640) eval loss: 1.2868
+- Final SFT step (640) eval perplexity: 3.6210
+
+The SFT process successfully aligned the model to follow prompt formats and drastically reduced perplexity on conversational validation targets.
+
+## Chat Format
+
+This model uses the standard **ChatML** system format. A typical chat turn looks like:
+
+```text
+<|im_start|>user
+Write a short email to my team about meeting tomorrow.<|im_end|>
+<|im_start|>assistant
+Subject: Meeting Tomorrow...<|im_end|>
+```
+
+## Intended Use
+
+- Assistant-style task completion
+- Multi-turn conversational chat
+- Zero-shot and few-shot instruction-following
+- Educational use and custom model inference experimentation
+
+## Out-of-Scope / Limitations
+
+- Safety-critical domains (medical, legal, financial advice)
+- Deployment in production without robust safety classifiers and filters
+- Handling long contexts beyond the 1,024-token limit
+- Language support beyond English (which dominates the smoltalk dataset)
+
+## Tokenization
+
+- Tokenizer: `tiktoken` with `cl100k_base` base ranks
+- Patched special tokens:
+  - `<|endoftext|>` = 100257 (EOS/PAD)
+  - `<|im_start|>` = 100264
+  - `<|im_end|>` = 100265
+  - `<|endofprompt|>` = 100276