Museko-125M final model with benchmark and trajectory charts

README.md

@@ -0,0 +1,66 @@
+---
+license: apache-2.0
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- llama
+- small-language-model
+- from-scratch
+datasets:
+- HuggingFaceFW/fineweb-edu
+---
+
+# Museko-125M
+
+Museko-125M is a 123M parameter language model I trained from scratch. It targets the sub-150M size class, and across the standard small-model benchmarks it currently comes out ahead of the other models in that range.
+
+It is a Llama style decoder pretrained on FineWeb-Edu with a small amount of synthetic arithmetic mixed in. The arithmetic mix is why it does noticeably better on math than most models its size while still holding solid general knowledge and reasoning.
+
+## Benchmarks
+
+All numbers are 0-shot, measured with lm-evaluation-harness and the ArithMark 2.0 benchmark.
+
+| Benchmark | Museko-125M | Previous best in class |
+|---|---|---|
+| HellaSwag | 33.2 | 30.6 |
+| ARC-Easy | 54.0 | 48.6 |
+| ARC-Challenge | 25.7 | 25.4 |
+| PIQA | 63.2 | 64.5 |
+| ArithMark 2.0 | 74.3 | 66.7 |
+
+![Benchmark comparison](bench.png)
+
+By the leaderboard averaging (ARC-Easy and ARC-Challenge combined) this comes to about 52.6, ahead of the previous top model at 49.71. It leads on four of the five benchmarks.
+
+![Board average during training](trajectory.png)
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+name = "TobiasLogic/Museko-125M"
+tok = AutoTokenizer.from_pretrained(name)
+model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=torch.bfloat16)
+
+prompt = "The moon is"
+ids = tok(prompt, return_tensors="pt")
+out = model.generate(**ids, max_new_tokens=40)
+print(tok.decode(out[0], skip_special_tokens=True))
+```
+
+## Training setup
+
+- Architecture: Llama, 12 layers, hidden size 768, 12 attention heads, 2048 context
+- Parameters: about 123M
+- Tokenizer: SmolLM2-135M (49,152 vocab)
+- Data: FineWeb-Edu plus roughly 10% synthetic arithmetic
+- Optimizer: AdamW, cosine schedule from 6e-4 down to 6e-5, 1,000 step warmup, weight decay 0.1
+- Effective batch: about 0.52M tokens per step
+
+## Notes
+
+Weights are open under Apache 2.0. This release is a from-scratch model and remains a strong general baseline for its size while being particularly good at arithmetic.

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "pad_token_id": null,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "use_cache": true,
+  "vocab_size": 49152
+}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.14.1",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:234232686967d0c93f66204f18dcf6c517d0666e83d4f7814d0fbe3e42ba210e
+size 490822600

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 8192,
+  "pad_token": null,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}