Switch weights to bf16, re-exported from the fp32 checkpoint

README.md

@@ -29,7 +29,7 @@ datasets:
 
 # KeyLM-75M-Instruct
 
-KeyLM-75M-Instruct is a 75M parameter instruction-tuned language model trained from scratch on approximately 18 billion tokens. That training budget is a small fraction of what comparable small models use (SmolLM-135M was trained on 600B tokens and SmolLM2-135M on 2T). Despite this, it is competitive on instruction following, outperforming SmolLM-135M-Instruct on IFEval while using about half the parameters and a fraction of the data.
+KeyLM-75M-Instruct is a 75M parameter instruction-tuned language model trained from scratch on approximately 18 billion tokens. That training budget is a small fraction of what comparable small models use (SmolLM-135M was trained on roughly 600B tokens, SmolLM2-135M on roughly 2T). Despite this, it is competitive on instruction following, outperforming SmolLM-135M-Instruct on IFEval while using about half the parameters and a fraction of the data.
 
 ## Table of Contents
 
@@ -53,13 +53,15 @@ KeyLM is a compact decoder-only transformer built on the standard small-model re
 | Attention heads | 8 (2 KV heads, GQA) |
 | Context length | 2048 |
 | Vocabulary | 12,020 (ByteLevel BPE) |
-| Precision | float16 |
+| Precision | bfloat16 |
 | Training tokens | ~18B |
 
 GGUF builds for `llama.cpp`, LM Studio, and Ollama are available at [KeyLM-75M-Instruct-GGUF](https://huggingface.co/Eclipse-Senpai/KeyLM-75M-Instruct-GGUF).
 
 ## How to Use
 
+KeyLM ships its own modeling code, so load it with `trust_remote_code=True`. It requires `transformers>=4.51`.
+
 ```python
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -67,7 +69,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model_id = "Eclipse-Senpai/KeyLM-75M-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    model_id, trust_remote_code=True, torch_dtype=torch.float16
+    model_id, trust_remote_code=True, torch_dtype=torch.bfloat16
 )
 
 messages = [{"role": "user", "content": "What is the capital of France?"}]
@@ -81,6 +83,8 @@ outputs = model.generate(
 print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True))
 ```
 
+The model uses a plain `User:` / `Assistant:` chat format, applied automatically by `apply_chat_template`. Assistant turns end with `</s>`.
+
 ## Evaluation
 
 ### Instruction following (IFEval)

config.json

@@ -9,11 +9,11 @@
   },
   "vocab_size": 12020,
   "hidden_size": 512,
-  "intermediate_size": 1280,
-  "num_hidden_layers": 24,
+  "head_dim": 64,
   "num_attention_heads": 8,
   "num_key_value_heads": 2,
-  "head_dim": 64,
+  "intermediate_size": 1280,
+  "num_hidden_layers": 24,
   "max_position_embeddings": 2048,
   "rope_theta": 10000.0,
   "rms_norm_eps": 1e-06,
@@ -26,5 +26,5 @@
   "bos_token_id": 1,
   "eos_token_id": 2,
   "pad_token_id": 2,
-  "torch_dtype": "float16"
+  "torch_dtype": "bfloat16"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:62a2f1202bf8c44f7839d0f402c81977d8516936a6a7aa70bc8cebd210791b4b
-size 150531664
+oid sha256:58ff80ee933def1960d75556c3770a016f56cb3f8f3891f01ca9501ca6cf17dc
+size 150531928

tokenizer_config.json

@@ -6,8 +6,8 @@
   "tokenizer_class": "PreTrainedTokenizerFast",
   "unk_token": "[UNK]",
   "vocab_size": 12020,
-  "chat_template": "{% for message in messages %}{% if message['role'] == 'user' %}{% if loop.index0 > 0 %}\n{% endif %}User: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}</s>{% endif %}{% endfor %}{% if add_generation_prompt %}Assistant: {% endif %}",
   "add_bos_token": false,
   "add_eos_token": false,
-  "clean_up_tokenization_spaces": false
+  "clean_up_tokenization_spaces": false,
+  "chat_template": "{% for message in messages %}{% if message['role'] == 'user' %}{% if loop.index0 > 0 %}\n{% endif %}User: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}</s>{% endif %}{% endfor %}{% if add_generation_prompt %}Assistant: {% endif %}"
 }