mecoffey commited on
Commit
10d8d08
·
1 Parent(s): 476420d

uploading tuned model

Browse files
README.md CHANGED
@@ -1,3 +1,44 @@
1
  ---
 
 
2
  license: apache-2.0
 
 
 
 
 
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ language:
3
+ - en
4
  license: apache-2.0
5
+ library_name: transformers
6
+ tags:
7
+ - gemma
8
+ - text-generation
9
+ - pytorch
10
+ - causal-lm
11
+ - custom-architecture
12
+ datasets:
13
+ - HuggingFaceTB/cosmopedia
14
+ pipeline_tag: text-generation
15
  ---
16
+
17
+ # The Original Model: GrownUpBaby-110M 👶➡️👨‍🎓
18
+
19
+ "Bedtime stories grew up."
20
+
21
+ **GrownUpBaby-110M** is a compact, Gemma-style causal LLM (110,304,256 parameters) trained from scratch to be a capable storyteller and creative assistant, with strong coherence and thematic control despite its size.
22
+ It's the grown-up counterpart to my earlier **[Exquisique/BabyLangModel](https://huggingface.co/Exquisique/BabyLangModel)**—an LLM with fewer parameters (30M) trained from scratch on TinyStories to generate short, simple narratives for young readers.
23
+ With more room to breathe, GrownUpBaby aims for richer voice, longer arcs, and cleaner pacing—built to read like a **master storyteller** on consumer hardware.
24
+
25
+ # The Fine-Tuned Model: NPC_brain
26
+
27
+ Using the Storytelling prowess of GrownUpBaby, I added training data to specialize the model in formulating NPC descriptions and backstories for an NPC Generator!
28
+
29
+ In most cases, I took a cool thing, and made it useless for 98% of huggingface. But think of the lore you can create!
30
+
31
+ ## 📚 Training Data
32
+ I trained this model using **Modal** ![alt text](image.png) and a custom dataset found at https://huggingface.co/datasets/mecoffey/npc_dataset
33
+
34
+
35
+ ## ⚠️ Limitations & Bias
36
+
37
+ * **Size:** At 110M parameters, this model has limited "world knowledge" compared to 7B+ models. It is best suited for creative writing and simple instruction following.
38
+ * **Hallucinations:** It may generate plausible-sounding but factually incorrect information.
39
+ * **Language:** Trained primarily on English educational and story data.
40
+
41
+ ## 👨‍💻 Author
42
+
43
+ Trained by **Exquisique**
44
+ Fine-Tuned by **mecoffey**
config.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "hidden_size": 768,
3
+ "num_layers": 12,
4
+ "intermediate_size": 3072,
5
+ "num_attention_heads": 12,
6
+ "num_key_value_heads": 4,
7
+ "max_position_embeddings": 1024,
8
+ "vocab_size": 50257,
9
+ "rope_theta": 10000.0,
10
+ "rms_norm_eps": 1e-06,
11
+ "hidden_act": "gelu_pytorch_tanh",
12
+ "attention_dropout": 0.0,
13
+ "hidden_dropout": 0.0,
14
+ "tie_word_embeddings": false,
15
+ "initializer_range": 0.02,
16
+ "architectures": [
17
+ "GemmaForCausalLM"
18
+ ],
19
+ "auto_map": {
20
+ "AutoConfig": "model.GemmaConfig",
21
+ "AutoModelForCausalLM": "model.GemmaForCausalLM"
22
+ },
23
+ "model_type": "gemma_custom",
24
+ "torch_dtype": "float32"
25
+ }
image.png ADDED
model.py ADDED
@@ -0,0 +1,540 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # =============================================================================
2
+ # Gemma-Style Model Architecture (110M Parameters)
3
+ # =============================================================================
4
+ """
5
+ This module implements a Gemma-style transformer architecture with:
6
+ - RMSNorm (Root Mean Square Layer Normalization)
7
+ - RoPE (Rotary Position Embeddings)
8
+ - GeGLU (Gated Linear Unit with GELU activation)
9
+ - GQA (Grouped Query Attention)
10
+
11
+ Architecture Specifications:
12
+ - hidden_size: 768
13
+ - num_layers: 12
14
+ - intermediate_size: 3072
15
+ - num_attention_heads: 12
16
+ - num_key_value_heads: 4 (GQA ratio of 3:1)
17
+ - vocab_size: 50257 (GPT-2 tokenizer)
18
+ """
19
+
20
+ import math
21
+ import logging
22
+ from typing import Optional, Tuple, Dict, Any, Union, List
23
+
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from torch import Tensor
28
+
29
+ from transformers import PretrainedConfig, PreTrainedModel
30
+ from transformers.modeling_outputs import CausalLMOutputWithPast
31
+
32
+ logger = logging.getLogger(__name__)
33
+
34
+
35
+ class GemmaConfig(PretrainedConfig):
36
+ """
37
+ Configuration class for the Gemma-style model.
38
+
39
+ Inherits from HuggingFace PretrainedConfig for compatibility with
40
+ AutoConfig and the transformers ecosystem.
41
+ """
42
+ model_type = "gemma_custom"
43
+
44
+ def __init__(
45
+ self,
46
+ hidden_size: int = 768,
47
+ num_layers: int = 12,
48
+ intermediate_size: int = 3072,
49
+ num_attention_heads: int = 12,
50
+ num_key_value_heads: int = 4,
51
+ max_position_embeddings: int = 1024,
52
+ vocab_size: int = 50257,
53
+ rope_theta: float = 10000.0,
54
+ rms_norm_eps: float = 1e-6,
55
+ hidden_act: str = "gelu_pytorch_tanh",
56
+ attention_dropout: float = 0.0,
57
+ hidden_dropout: float = 0.0,
58
+ tie_word_embeddings: bool = True,
59
+ initializer_range: float = 0.02,
60
+ **kwargs
61
+ ):
62
+ self.hidden_size = hidden_size
63
+ self.num_layers = num_layers
64
+ self.intermediate_size = intermediate_size
65
+ self.num_attention_heads = num_attention_heads
66
+ self.num_key_value_heads = num_key_value_heads
67
+ self.max_position_embeddings = max_position_embeddings
68
+ self.vocab_size = vocab_size
69
+ self.rope_theta = rope_theta
70
+ self.rms_norm_eps = rms_norm_eps
71
+ self.hidden_act = hidden_act
72
+ self.attention_dropout = attention_dropout
73
+ self.hidden_dropout = hidden_dropout
74
+ self.initializer_range = initializer_range
75
+
76
+ # Derived attributes
77
+ self.head_dim = self.hidden_size // self.num_attention_heads
78
+ self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
79
+
80
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
81
+
82
+
83
+ class RMSNorm(nn.Module):
84
+ """
85
+ Root Mean Square Layer Normalization.
86
+
87
+ RMSNorm(x) = x * rsqrt(mean(x^2) + eps) * weight
88
+ """
89
+
90
+ def __init__(self, hidden_size: int, eps: float = 1e-6):
91
+ super().__init__()
92
+ self.weight = nn.Parameter(torch.ones(hidden_size))
93
+ self.variance_epsilon = eps
94
+
95
+ def forward(self, hidden_states: Tensor) -> Tensor:
96
+ input_dtype = hidden_states.dtype
97
+ hidden_states = hidden_states.to(torch.float32)
98
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
99
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
100
+ return self.weight * hidden_states.to(input_dtype)
101
+
102
+
103
+ class RotaryEmbedding(nn.Module):
104
+ """Rotary Position Embedding (RoPE)."""
105
+
106
+ def __init__(self, dim: int, max_position_embeddings: int = 1024, theta: float = 10000.0):
107
+ super().__init__()
108
+ self.dim = dim
109
+ self.max_position_embeddings = max_position_embeddings
110
+ self.theta = theta
111
+
112
+ inv_freq = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
113
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
114
+ self._set_cos_sin_cache(max_position_embeddings)
115
+
116
+ def _set_cos_sin_cache(self, seq_len: int):
117
+ t = torch.arange(seq_len, dtype=torch.float32)
118
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
119
+ emb = torch.cat((freqs, freqs), dim=-1)
120
+ self.register_buffer("cos_cached", emb.cos(), persistent=False)
121
+ self.register_buffer("sin_cached", emb.sin(), persistent=False)
122
+
123
+ def forward(self, x: Tensor, position_ids: Tensor) -> Tuple[Tensor, Tensor]:
124
+ seq_len = position_ids.max() + 1
125
+ if seq_len > self.cos_cached.shape[0]:
126
+ self._set_cos_sin_cache(seq_len)
127
+ self.cos_cached = self.cos_cached.to(x.device)
128
+ self.sin_cached = self.sin_cached.to(x.device)
129
+
130
+ cos = self.cos_cached[position_ids].unsqueeze(2)
131
+ sin = self.sin_cached[position_ids].unsqueeze(2)
132
+ return cos.to(x.dtype), sin.to(x.dtype)
133
+
134
+
135
+ def rotate_half(x: Tensor) -> Tensor:
136
+ x1 = x[..., : x.shape[-1] // 2]
137
+ x2 = x[..., x.shape[-1] // 2 :]
138
+ return torch.cat((-x2, x1), dim=-1)
139
+
140
+
141
+ def apply_rotary_pos_emb(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> Tuple[Tensor, Tensor]:
142
+ q_embed = (q * cos) + (rotate_half(q) * sin)
143
+ k_embed = (k * cos) + (rotate_half(k) * sin)
144
+ return q_embed, k_embed
145
+
146
+
147
+ class GemmaMLP(nn.Module):
148
+ """Gemma-style MLP with GeGLU activation."""
149
+
150
+ def __init__(self, config: GemmaConfig):
151
+ super().__init__()
152
+ self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
153
+ self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
154
+ self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
155
+ self.act_fn = nn.GELU(approximate="tanh")
156
+
157
+ def forward(self, x: Tensor) -> Tensor:
158
+ gate = self.act_fn(self.gate_proj(x))
159
+ up = self.up_proj(x)
160
+ return self.down_proj(gate * up)
161
+
162
+
163
+ class GemmaAttention(nn.Module):
164
+ """Grouped Query Attention (GQA) module."""
165
+
166
+ def __init__(self, config: GemmaConfig, layer_idx: int):
167
+ super().__init__()
168
+ self.config = config
169
+ self.layer_idx = layer_idx
170
+ self.hidden_size = config.hidden_size
171
+ self.num_heads = config.num_attention_heads
172
+ self.head_dim = config.head_dim
173
+ self.num_key_value_heads = config.num_key_value_heads
174
+ self.num_key_value_groups = config.num_key_value_groups
175
+ self.attention_dropout = config.attention_dropout
176
+
177
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
178
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
179
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
180
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
181
+
182
+ self.rotary_emb = RotaryEmbedding(
183
+ self.head_dim,
184
+ max_position_embeddings=config.max_position_embeddings,
185
+ theta=config.rope_theta
186
+ )
187
+
188
+ def forward(
189
+ self,
190
+ hidden_states: Tensor,
191
+ attention_mask: Optional[Tensor] = None,
192
+ position_ids: Optional[Tensor] = None,
193
+ ) -> Tensor:
194
+ batch_size, seq_len, _ = hidden_states.shape
195
+
196
+ query_states = self.q_proj(hidden_states)
197
+ key_states = self.k_proj(hidden_states)
198
+ value_states = self.v_proj(hidden_states)
199
+
200
+ query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim)
201
+ key_states = key_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
202
+ value_states = value_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
203
+
204
+ cos, sin = self.rotary_emb(query_states, position_ids)
205
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
206
+
207
+ query_states = query_states.transpose(1, 2)
208
+ key_states = key_states.transpose(1, 2)
209
+ value_states = value_states.transpose(1, 2)
210
+
211
+ if self.num_key_value_groups > 1:
212
+ key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
213
+ value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
214
+
215
+ scale = 1.0 / math.sqrt(self.head_dim)
216
+ attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * scale
217
+
218
+ if attention_mask is not None:
219
+ attn_weights = attn_weights + attention_mask
220
+
221
+ attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
222
+ attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
223
+
224
+ attn_output = torch.matmul(attn_weights, value_states)
225
+ attn_output = attn_output.transpose(1, 2).contiguous()
226
+ attn_output = attn_output.reshape(batch_size, seq_len, self.hidden_size)
227
+
228
+ return self.o_proj(attn_output)
229
+
230
+
231
+ class GemmaDecoderLayer(nn.Module):
232
+ """Single transformer decoder layer with pre-norm architecture."""
233
+
234
+ def __init__(self, config: GemmaConfig, layer_idx: int):
235
+ super().__init__()
236
+ self.hidden_size = config.hidden_size
237
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
238
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
239
+ self.self_attn = GemmaAttention(config, layer_idx)
240
+ self.mlp = GemmaMLP(config)
241
+
242
+ def forward(
243
+ self,
244
+ hidden_states: Tensor,
245
+ attention_mask: Optional[Tensor] = None,
246
+ position_ids: Optional[Tensor] = None,
247
+ ) -> Tensor:
248
+ residual = hidden_states
249
+ hidden_states = self.input_layernorm(hidden_states)
250
+ hidden_states = self.self_attn(hidden_states, attention_mask, position_ids)
251
+ hidden_states = residual + hidden_states
252
+
253
+ residual = hidden_states
254
+ hidden_states = self.post_attention_layernorm(hidden_states)
255
+ hidden_states = self.mlp(hidden_states)
256
+ hidden_states = residual + hidden_states
257
+
258
+ return hidden_states
259
+
260
+
261
+ class GemmaModel(nn.Module):
262
+ """Core Gemma transformer model (without LM head)."""
263
+
264
+ def __init__(self, config: GemmaConfig):
265
+ super().__init__()
266
+ self.config = config
267
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
268
+ self.layers = nn.ModuleList([
269
+ GemmaDecoderLayer(config, layer_idx)
270
+ for layer_idx in range(config.num_layers)
271
+ ])
272
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
273
+ self.gradient_checkpointing = False
274
+
275
+ def forward(
276
+ self,
277
+ input_ids: Tensor,
278
+ attention_mask: Optional[Tensor] = None,
279
+ position_ids: Optional[Tensor] = None,
280
+ ) -> Tensor:
281
+ batch_size, seq_len = input_ids.shape
282
+ hidden_states = self.embed_tokens(input_ids)
283
+
284
+ if position_ids is None:
285
+ position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
286
+
287
+ causal_mask = self._create_causal_mask(seq_len, hidden_states.device, hidden_states.dtype)
288
+
289
+ for layer in self.layers:
290
+ if self.gradient_checkpointing and self.training:
291
+ hidden_states = torch.utils.checkpoint.checkpoint(
292
+ layer, hidden_states, causal_mask, position_ids, use_reentrant=False
293
+ )
294
+ else:
295
+ hidden_states = layer(hidden_states, causal_mask, position_ids)
296
+
297
+ return self.norm(hidden_states)
298
+
299
+ def _create_causal_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tensor:
300
+ mask = torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype)
301
+ mask = torch.triu(mask, diagonal=1)
302
+ return mask.unsqueeze(0).unsqueeze(0)
303
+
304
+
305
+ class GemmaForCausalLM(PreTrainedModel):
306
+ """
307
+ Gemma model with language modeling head for causal text generation.
308
+
309
+ Inherits from HuggingFace PreTrainedModel for full compatibility with
310
+ AutoModelForCausalLM and the transformers ecosystem.
311
+ """
312
+ config_class = GemmaConfig
313
+ supports_gradient_checkpointing = True
314
+ _no_split_modules = ["GemmaDecoderLayer"]
315
+ _supports_param_buffer_assignment = False # Fix for accelerate weight loading
316
+
317
+ def __init__(self, config: GemmaConfig):
318
+ super().__init__(config)
319
+ self.config = config
320
+
321
+ self.model = GemmaModel(config)
322
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
323
+
324
+ if config.tie_word_embeddings:
325
+ self.lm_head.weight = self.model.embed_tokens.weight
326
+
327
+ self.post_init()
328
+
329
+ @classmethod
330
+ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
331
+ """
332
+ Custom from_pretrained that properly loads weights for this custom model.
333
+ This overrides the default behavior to ensure weights are loaded correctly.
334
+ Supports both .safetensors and .bin formats.
335
+ """
336
+ import os
337
+ from huggingface_hub import hf_hub_download
338
+
339
+ # Get config
340
+ trust_remote_code = kwargs.pop("trust_remote_code", True)
341
+ torch_dtype = kwargs.pop("torch_dtype", None)
342
+ device_map = kwargs.pop("device_map", None)
343
+
344
+ # Load config
345
+ config = cls.config_class.from_pretrained(
346
+ pretrained_model_name_or_path,
347
+ trust_remote_code=trust_remote_code,
348
+ **kwargs
349
+ )
350
+
351
+ # Create model
352
+ model = cls(config)
353
+
354
+ # Find weight file (try safetensors first, then pytorch_model.bin)
355
+ weight_file = None
356
+ if os.path.isdir(pretrained_model_name_or_path):
357
+ # Check local directory
358
+ safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors")
359
+ bin_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
360
+
361
+ if os.path.exists(safetensors_path):
362
+ weight_file = safetensors_path
363
+ elif os.path.exists(bin_path):
364
+ weight_file = bin_path
365
+ else:
366
+ logger.warning(f"No weight file found in {pretrained_model_name_or_path}")
367
+ return model
368
+ else:
369
+ # Download from hub (try safetensors first)
370
+ try:
371
+ weight_file = hf_hub_download(
372
+ repo_id=pretrained_model_name_or_path,
373
+ filename="model.safetensors"
374
+ )
375
+ except Exception:
376
+ weight_file = hf_hub_download(
377
+ repo_id=pretrained_model_name_or_path,
378
+ filename="pytorch_model.bin"
379
+ )
380
+
381
+ # Load weights based on file format
382
+ if weight_file.endswith(".safetensors"):
383
+ try:
384
+ from safetensors.torch import load_file
385
+ state_dict = load_file(weight_file)
386
+ except ImportError:
387
+ logger.warning("safetensors not installed, trying torch.load instead")
388
+ state_dict = torch.load(weight_file, map_location="cpu")
389
+ else:
390
+ state_dict = torch.load(weight_file, map_location="cpu")
391
+
392
+ model.load_state_dict(state_dict, strict=False)
393
+
394
+ # Handle dtype and device
395
+ if torch_dtype is not None:
396
+ model = model.to(torch_dtype)
397
+ if device_map == "auto":
398
+ if torch.cuda.is_available():
399
+ model = model.to("cuda")
400
+ elif device_map is not None:
401
+ model = model.to(device_map)
402
+
403
+ return model
404
+
405
+ def get_input_embeddings(self):
406
+ return self.model.embed_tokens
407
+
408
+ def set_input_embeddings(self, value):
409
+ self.model.embed_tokens = value
410
+
411
+ def get_output_embeddings(self):
412
+ return self.lm_head
413
+
414
+ def set_output_embeddings(self, new_embeddings):
415
+ self.lm_head = new_embeddings
416
+
417
+ def _init_weights(self, module: nn.Module):
418
+ std = self.config.initializer_range
419
+ if isinstance(module, nn.Linear):
420
+ module.weight.data.normal_(mean=0.0, std=std)
421
+ if module.bias is not None:
422
+ module.bias.data.zero_()
423
+ elif isinstance(module, nn.Embedding):
424
+ module.weight.data.normal_(mean=0.0, std=std)
425
+
426
+ def num_parameters(self, only_trainable: bool = False) -> int:
427
+ return sum(p.numel() for p in self.parameters() if not only_trainable or p.requires_grad)
428
+
429
+ def forward(
430
+ self,
431
+ input_ids: Tensor,
432
+ attention_mask: Optional[Tensor] = None,
433
+ position_ids: Optional[Tensor] = None,
434
+ labels: Optional[Tensor] = None,
435
+ return_dict: Optional[bool] = None,
436
+ **kwargs
437
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
438
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
439
+
440
+ hidden_states = self.model(input_ids, attention_mask, position_ids)
441
+ logits = self.lm_head(hidden_states)
442
+
443
+ loss = None
444
+ if labels is not None:
445
+ shift_logits = logits[..., :-1, :].contiguous()
446
+ shift_labels = labels[..., 1:].contiguous()
447
+ loss_fct = nn.CrossEntropyLoss()
448
+ loss = loss_fct(
449
+ shift_logits.view(-1, self.config.vocab_size),
450
+ shift_labels.view(-1)
451
+ )
452
+
453
+ if not return_dict:
454
+ output = (logits,)
455
+ return ((loss,) + output) if loss is not None else output
456
+
457
+ return CausalLMOutputWithPast(
458
+ loss=loss,
459
+ logits=logits,
460
+ past_key_values=None,
461
+ hidden_states=None,
462
+ attentions=None,
463
+ )
464
+
465
+ def prepare_inputs_for_generation(
466
+ self, input_ids, past_key_values=None, attention_mask=None, **kwargs
467
+ ):
468
+ if past_key_values is not None:
469
+ input_ids = input_ids[:, -1:]
470
+
471
+ position_ids = kwargs.get("position_ids", None)
472
+ if attention_mask is not None and position_ids is None:
473
+ position_ids = attention_mask.long().cumsum(-1) - 1
474
+ position_ids.masked_fill_(attention_mask == 0, 1)
475
+ if past_key_values:
476
+ position_ids = position_ids[:, -1].unsqueeze(-1)
477
+
478
+ return {
479
+ "input_ids": input_ids,
480
+ "attention_mask": attention_mask,
481
+ "position_ids": position_ids,
482
+ "past_key_values": past_key_values,
483
+ }
484
+
485
+ def generate(
486
+ self,
487
+ input_ids: Tensor,
488
+ max_new_tokens: int = 50,
489
+ temperature: float = 1.0,
490
+ top_k: int = 50,
491
+ top_p: float = 0.95,
492
+ do_sample: bool = True,
493
+ eos_token_id: Optional[int] = None,
494
+ **kwargs
495
+ ) -> Tensor:
496
+ """Custom generate method for simple autoregressive generation."""
497
+ self.eval()
498
+
499
+ for _ in range(max_new_tokens):
500
+ with torch.no_grad():
501
+ outputs = self.forward(input_ids)
502
+ next_token_logits = outputs.logits[:, -1, :] / temperature
503
+
504
+ if top_k > 0:
505
+ indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
506
+ next_token_logits[indices_to_remove] = float("-inf")
507
+
508
+ if top_p < 1.0:
509
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
510
+ cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
511
+ sorted_indices_to_remove = cumulative_probs > top_p
512
+ sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
513
+ sorted_indices_to_remove[..., 0] = 0
514
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
515
+ next_token_logits[indices_to_remove] = float("-inf")
516
+
517
+ if do_sample:
518
+ probs = F.softmax(next_token_logits, dim=-1)
519
+ next_token = torch.multinomial(probs, num_samples=1)
520
+ else:
521
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
522
+
523
+ input_ids = torch.cat([input_ids, next_token], dim=-1)
524
+
525
+ if eos_token_id is not None and (next_token == eos_token_id).all():
526
+ break
527
+
528
+ return input_ids
529
+
530
+
531
+ def create_model_from_config(config_dict: Optional[Dict[str, Any]] = None) -> GemmaForCausalLM:
532
+ """Factory function to create a model from a configuration dictionary."""
533
+ if config_dict is None:
534
+ config = GemmaConfig()
535
+ else:
536
+ config = GemmaConfig(**{k: v for k, v in config_dict.items() if hasattr(GemmaConfig, k) or k in ['hidden_size', 'num_layers', 'intermediate_size', 'num_attention_heads', 'num_key_value_heads', 'max_position_embeddings', 'vocab_size', 'rope_theta', 'rms_norm_eps', 'hidden_act', 'attention_dropout', 'hidden_dropout', 'tie_word_embeddings', 'initializer_range']})
537
+
538
+ model = GemmaForCausalLM(config)
539
+ logger.info(f"Created model with {model.num_parameters():,} parameters")
540
+ return model
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:35cc7233b55b81823e54b1b752418618c88f7117dfca30896ab8dbfcaabc57c5
3
+ size 362058272
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ datasets
4
+ safetensors
5
+ accelerate
6
+ modal
7
+ pathlib
8
+ gradio
9
+ huggingface-hub
special_tokens_map.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": "<|endoftext|>",
3
+ "eos_token": "<|endoftext|>",
4
+ "unk_token": "<|endoftext|>",
5
+ "pad_token": "<|endoftext|>"
6
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "backend": "tokenizers",
4
+ "bos_token": "<|endoftext|>",
5
+ "eos_token": "<|endoftext|>",
6
+ "errors": "replace",
7
+ "is_local": false,
8
+ "model_max_length": 1024,
9
+ "pad_token": null,
10
+ "tokenizer_class": "GPT2Tokenizer",
11
+ "unk_token": "<|endoftext|>"
12
+ }