GautierP commited on
Commit
abb24a4
·
verified ·
1 Parent(s): 51c2d36

Upload folder using huggingface_hub

Browse files
README.md ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language: en
3
+ library_name: transformers
4
+ pipeline_tag: text-generation
5
+ tags:
6
+ - mini-llm
7
+ - from-scratch
8
+ - small-language-model
9
+ ---
10
+
11
+ # Mini-LLM Assistant 70M
12
+
13
+ A ~70M parameter conversational assistant trained **entirely from scratch** with
14
+ PyTorch — no pretrained weights, no distillation. Tokenizer, pretraining and
15
+ supervised fine-tuning all belong to the
16
+ [Mini-LLM](https://github.com/Mini-llm-lab/Mini-LLM) project.
17
+
18
+ This repository is the SFT assistant derived from the `mini-llm-base-70M-22000steps.pt`
19
+ base checkpoint (not published).
20
+
21
+ ## Model details
22
+
23
+ | | |
24
+ | -------------- | ------------------------------------------------------------------------------------------------------------------------------- |
25
+ | Parameters | 69,975,680 (input embedding shared with the output head) |
26
+ | Architecture | Llama-style decoder: RMSNorm, RoPE, causal attention, SwiGLU |
27
+ | Context length | 512 tokens |
28
+ | Vocabulary | 16,384 (BPE) |
29
+ | Pretraining | ~1.44B tokens (FineWeb-Edu + Wikipedia), 22,000 steps — base val loss 3.030 |
30
+ | Fine-tuning | 2 epochs of SFT on Dolly-15k, No Robots, OpenAssistant oasst2, alpaca-cleaned and identity/small-talk examples — val loss 2.718 |
31
+
32
+ ## Usage
33
+
34
+ ```python
35
+ from transformers import AutoModelForCausalLM, AutoTokenizer
36
+
37
+ repo = "GautierP/MiniLLM-70M"
38
+ tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
39
+ model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
40
+
41
+ messages = [{"role": "user", "content": "What is the capital of France?"}]
42
+ inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
43
+ out = model.generate(inputs, max_new_tokens=128)
44
+ print(tok.decode(out[0][inputs.shape[1]:], skip_special_tokens=True))
45
+ ```
46
+
47
+ `trust_remote_code=True` is required — the architecture ships with the repository
48
+ (`modeling_mini_llm.py`, `configuration_mini_llm.py`). Sampling defaults
49
+ (temperature 0.5, top_k 40, repetition penalty 1.15) live in `generation_config.json`.
50
+
51
+ The training format is single-turn and fixed; `chat_template.jinja` implements it
52
+ exactly:
53
+
54
+ ```
55
+ User: <question>
56
+ Assistant: <answer>
57
+ ```
58
+
59
+ Multi-turn history is never sent to the model: it anchors on its own previous answers
60
+ and degrades when it sees them, so each question is answered independently.
61
+
62
+ ## Limitations
63
+
64
+ Very small model trained on a modest budget: answers are often fluent and in the right
65
+ format, but **factual errors and hallucinations are common**. Do not use it for factual
66
+ lookup or any production task. English only.
67
+
68
+ ## Links
69
+
70
+ - Code, training pipeline and full documentation: https://github.com/Mini-llm-lab/Mini-LLM
chat_template.jinja ADDED
@@ -0,0 +1 @@
 
 
1
+ {% for m in messages %}{{ 'User: ' + m['content'].strip() + '\n' if m['role'] == 'user' else 'Assistant: ' + m['content'].strip() + '\n' }}{% endfor %}{% if add_generation_prompt %}Assistant:{% endif %}
config.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "MiniLlmForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "configuration_mini_llm.MiniLlmConfig",
7
+ "AutoModelForCausalLM": "modeling_mini_llm.MiniLlmForCausalLM"
8
+ },
9
+ "dtype": "float32",
10
+ "eos_token_id": 0,
11
+ "hidden_size": 640,
12
+ "intermediate_size": 1728,
13
+ "max_position_embeddings": 512,
14
+ "model_type": "mini_llm",
15
+ "num_attention_heads": 10,
16
+ "num_hidden_layers": 12,
17
+ "pad_token_id": 0,
18
+ "rms_norm_eps": 1e-05,
19
+ "rope_theta": 10000.0,
20
+ "tie_word_embeddings": true,
21
+ "transformers_version": "5.19.0",
22
+ "vocab_size": 16384
23
+ }
configuration_mini_llm.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class MiniLlmConfig(PretrainedConfig):
5
+ model_type = "mini_llm"
6
+
7
+ def __init__(
8
+ self,
9
+ vocab_size=16384,
10
+ hidden_size=512,
11
+ num_hidden_layers=8,
12
+ num_attention_heads=8,
13
+ intermediate_size=1408,
14
+ max_position_embeddings=512,
15
+ rms_norm_eps=1e-5,
16
+ rope_theta=10000.0,
17
+ tie_word_embeddings=True,
18
+ **kwargs,
19
+ ):
20
+ super().__init__(
21
+ tie_word_embeddings=tie_word_embeddings,
22
+ eos_token_id=kwargs.pop("eos_token_id", 0),
23
+ pad_token_id=kwargs.pop("pad_token_id", 0),
24
+ **kwargs,
25
+ )
26
+ self.vocab_size = vocab_size
27
+ self.hidden_size = hidden_size
28
+ self.num_hidden_layers = num_hidden_layers
29
+ self.num_attention_heads = num_attention_heads
30
+ self.intermediate_size = intermediate_size
31
+ self.max_position_embeddings = max_position_embeddings
32
+ self.rms_norm_eps = rms_norm_eps
33
+ self.rope_theta = rope_theta
34
+
35
+ @classmethod
36
+ def from_cfg(cls, cfg, **kwargs):
37
+ return cls(
38
+ vocab_size=cfg.vocab,
39
+ hidden_size=cfg.dim,
40
+ num_hidden_layers=cfg.layers,
41
+ num_attention_heads=cfg.heads,
42
+ intermediate_size=cfg.ffn,
43
+ max_position_embeddings=cfg.ctx,
44
+ **kwargs,
45
+ )
generation_config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "do_sample": true,
3
+ "eos_token_id": 0,
4
+ "max_new_tokens": 256,
5
+ "pad_token_id": 0,
6
+ "repetition_penalty": 1.15,
7
+ "temperature": 0.5,
8
+ "top_k": 40,
9
+ "transformers_version": "5.19.0"
10
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b9102687b74b4ad4ec77c05a6d8c968fdeb58dea8ed3656b701b617e10fb6348
3
+ size 280041904
modeling_mini_llm.py ADDED
@@ -0,0 +1,151 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ from transformers import PreTrainedModel
7
+ from transformers.generation import GenerationMixin
8
+ from transformers.modeling_outputs import CausalLMOutputWithPast
9
+
10
+ try:
11
+ from .configuration_mini_llm import MiniLlmConfig
12
+ except ImportError:
13
+ from configuration_mini_llm import MiniLlmConfig
14
+
15
+
16
+ class MiniLlmRMSNorm(nn.Module):
17
+ def __init__(self, d, eps=1e-5):
18
+ super().__init__()
19
+ self.w = nn.Parameter(torch.ones(d))
20
+ self.eps = eps
21
+
22
+ def forward(self, x):
23
+ n = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
24
+ return (x.float() * n).type_as(x) * self.w
25
+
26
+
27
+ def mini_rope_cache(T, hd, device, theta=10000.0):
28
+ inv = 1.0 / (theta ** (torch.arange(0, hd, 2, device=device).float() / hd))
29
+ f = torch.outer(torch.arange(T, device=device).float(), inv)
30
+ return f.cos(), f.sin()
31
+
32
+
33
+ def mini_apply_rope(x, cos, sin):
34
+ x1, x2 = x[..., ::2], x[..., 1::2]
35
+ cos, sin = cos[: x.size(2)], sin[: x.size(2)]
36
+ out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
37
+ return out.flatten(-2).type_as(x)
38
+
39
+
40
+ class MiniLlmAttention(nn.Module):
41
+ def __init__(self, config: MiniLlmConfig):
42
+ super().__init__()
43
+ d = config.hidden_size
44
+ self.h = config.num_attention_heads
45
+ self.qkv = nn.Linear(d, 3 * d, bias=False)
46
+ self.o = nn.Linear(d, d, bias=False)
47
+
48
+ def forward(self, x, cos, sin):
49
+ B, T, C = x.shape
50
+ q, k, v = self.qkv(x).view(B, T, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
51
+ q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
52
+ q, k = mini_apply_rope(q, cos, sin), mini_apply_rope(k, cos, sin)
53
+ y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
54
+ return self.o(y.transpose(1, 2).reshape(B, T, C))
55
+
56
+
57
+ class MiniLlmSwiGLU(nn.Module):
58
+ def __init__(self, config: MiniLlmConfig):
59
+ super().__init__()
60
+ d, ffn = config.hidden_size, config.intermediate_size
61
+ self.w1 = nn.Linear(d, ffn, bias=False)
62
+ self.w3 = nn.Linear(d, ffn, bias=False)
63
+ self.w2 = nn.Linear(ffn, d, bias=False)
64
+
65
+ def forward(self, x):
66
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
67
+
68
+
69
+ class MiniLlmBlock(nn.Module):
70
+ def __init__(self, config: MiniLlmConfig):
71
+ super().__init__()
72
+ self.n1, self.n2 = MiniLlmRMSNorm(config.hidden_size, config.rms_norm_eps), MiniLlmRMSNorm(
73
+ config.hidden_size, config.rms_norm_eps
74
+ )
75
+ self.attn, self.mlp = MiniLlmAttention(config), MiniLlmSwiGLU(config)
76
+
77
+ def forward(self, x, cos, sin):
78
+ x = x + self.attn(self.n1(x), cos, sin)
79
+ return x + self.mlp(self.n2(x))
80
+
81
+
82
+ class MiniLlmForCausalLM(PreTrainedModel, GenerationMixin):
83
+ config_class = MiniLlmConfig
84
+ base_model_prefix = "mini_llm"
85
+ supports_gradient_checkpointing = False
86
+ _tied_weights_keys = {"head.weight": "emb.weight"}
87
+
88
+ def __init__(self, config: MiniLlmConfig):
89
+ super().__init__(config)
90
+ self.config = config
91
+ self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
92
+ self.blocks = nn.ModuleList(MiniLlmBlock(config) for _ in range(config.num_hidden_layers))
93
+ self.norm = MiniLlmRMSNorm(config.hidden_size, config.rms_norm_eps)
94
+ self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
95
+ self.head.weight = self.emb.weight
96
+ cos, sin = mini_rope_cache(
97
+ config.max_position_embeddings,
98
+ config.hidden_size // config.num_attention_heads,
99
+ "cpu",
100
+ config.rope_theta,
101
+ )
102
+ self.register_buffer("cos", cos, persistent=True)
103
+ self.register_buffer("sin", sin, persistent=True)
104
+ self.post_init()
105
+ for n, p in self.named_parameters():
106
+ if n.endswith("attn.o.weight") or n.endswith("mlp.w2.weight"):
107
+ nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * config.num_hidden_layers))
108
+
109
+ def _init_weights(self, module):
110
+ if isinstance(module, (nn.Linear, nn.Embedding)):
111
+ nn.init.normal_(module.weight, 0.0, 0.02)
112
+
113
+ def get_input_embeddings(self):
114
+ return self.emb
115
+
116
+ def set_input_embeddings(self, value):
117
+ self.emb = value
118
+
119
+ def get_output_embeddings(self):
120
+ return self.head
121
+
122
+ def set_output_embeddings(self, value):
123
+ self.head = value
124
+
125
+ def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
126
+ kwargs.pop("past_key_values", None)
127
+ kwargs.pop("use_cache", None)
128
+ kwargs.pop("position_ids", None)
129
+ kwargs.pop("inputs_embeds", None)
130
+ if input_ids is None:
131
+ raise ValueError("input_ids is required")
132
+ x = self.emb(input_ids)
133
+ for b in self.blocks:
134
+ x = b(x, self.cos, self.sin)
135
+ logits = self.head(self.norm(x))
136
+ loss = None
137
+ if labels is not None:
138
+ loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), labels.view(-1))
139
+ return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
140
+
141
+ def prepare_inputs_for_generation(self, input_ids, **kwargs):
142
+ return {"input_ids": input_ids}
143
+
144
+
145
+ try:
146
+ from transformers import AutoConfig, AutoModelForCausalLM
147
+
148
+ AutoConfig.register("mini_llm", MiniLlmConfig)
149
+ AutoModelForCausalLM.register(MiniLlmConfig, MiniLlmForCausalLM)
150
+ except Exception:
151
+ pass
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": null,
4
+ "eos_token": "<eos>",
5
+ "model_max_length": 1000000000000000019884624838656,
6
+ "pad_token": "<eos>",
7
+ "tokenizer_class": "TokenizersBackend",
8
+ "unk_token": null
9
+ }