thebajajra commited on
Commit
96733b5
·
verified ·
1 Parent(s): 08a68ad

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .modeling_gemma3_biencoder import Gemma3EncoderForMaskedLM
config.json ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_sliding_window_pattern": 6,
3
+ "architectures": [
4
+ "Gemma3EncoderForMaskedLM"
5
+ ],
6
+ "attention_bias": false,
7
+ "attention_dropout": 0.0,
8
+ "attn_logit_softcapping": null,
9
+ "auto_map": {
10
+ "AutoModelForMaskedLM": "modeling_gemma3_biencoder.Gemma3EncoderForMaskedLM"
11
+ },
12
+ "bos_token_id": 2,
13
+ "dtype": "float32",
14
+ "eos_token_id": 1,
15
+ "final_logit_softcapping": null,
16
+ "head_dim": 256,
17
+ "hidden_activation": "gelu_pytorch_tanh",
18
+ "hidden_size": 640,
19
+ "initializer_range": 0.02,
20
+ "intermediate_size": 2048,
21
+ "layer_types": [
22
+ "sliding_attention",
23
+ "sliding_attention",
24
+ "sliding_attention",
25
+ "sliding_attention",
26
+ "sliding_attention",
27
+ "full_attention",
28
+ "sliding_attention",
29
+ "sliding_attention",
30
+ "sliding_attention",
31
+ "sliding_attention",
32
+ "sliding_attention",
33
+ "full_attention",
34
+ "sliding_attention",
35
+ "sliding_attention",
36
+ "sliding_attention",
37
+ "sliding_attention",
38
+ "sliding_attention",
39
+ "full_attention"
40
+ ],
41
+ "max_position_embeddings": 32768,
42
+ "model_type": "gemma3_text",
43
+ "num_attention_heads": 4,
44
+ "num_hidden_layers": 18,
45
+ "num_key_value_heads": 1,
46
+ "pad_token_id": 0,
47
+ "query_pre_attn_scalar": 256,
48
+ "rms_norm_eps": 1e-06,
49
+ "rope_local_base_freq": 10000.0,
50
+ "rope_scaling": null,
51
+ "rope_theta": 1000000.0,
52
+ "sliding_window": 512,
53
+ "transformers_version": "4.57.3",
54
+ "use_bidirectional_attention": true,
55
+ "use_cache": false,
56
+ "vocab_size": 262145
57
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e9dc0cc7558bc128b11280fbdbacf630a260a637110ad69d3de2f03ca9650093
3
+ size 1072422288
modeling_gemma3_biencoder.py ADDED
@@ -0,0 +1,282 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # gemma3_biencoder.py
2
+ from __future__ import annotations
3
+ import copy
4
+ import torch
5
+ import torch.nn as nn
6
+ from typing import Optional, Tuple, Union
7
+ from transformers.modeling_outputs import MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput
8
+ from transformers.models.gemma3.configuration_gemma3 import Gemma3TextConfig
9
+ from transformers.models.gemma3.modeling_gemma3 import (
10
+ Gemma3PreTrainedModel,
11
+ Gemma3TextModel,
12
+ )
13
+
14
+ class Gemma3EncoderForMaskedLM(Gemma3PreTrainedModel):
15
+ config_class = Gemma3TextConfig
16
+ base_model_prefix = "encoder"
17
+ _tied_weights_keys = ["lm_head.weight"]
18
+ _keys_to_ignore_on_load_missing = [r"lm_head\.weight"]
19
+
20
+ def __init__(self, config: Gemma3TextConfig):
21
+ cfg = copy.deepcopy(config)
22
+ if hasattr(cfg, "use_bidirectional_attention"):
23
+ cfg.use_bidirectional_attention = True
24
+ cfg.use_cache = False
25
+ super().__init__(cfg)
26
+
27
+ self.encoder = Gemma3TextModel(cfg)
28
+ self.vocab_size = cfg.vocab_size
29
+ self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
30
+ self.post_init() # calls tie_weights()
31
+
32
+ # Embeddings / head
33
+ def get_input_embeddings(self):
34
+ return self.encoder.embed_tokens
35
+
36
+ def set_input_embeddings(self, new_embeddings):
37
+ self.encoder.embed_tokens = new_embeddings
38
+
39
+ def get_output_embeddings(self):
40
+ return self.lm_head
41
+
42
+ def set_output_embeddings(self, new_head: nn.Module):
43
+ self.lm_head = new_head
44
+
45
+ # Keep vocab_size in sync; ensure pointer-tying
46
+ def tie_weights(self):
47
+ if hasattr(self.config, "vocab_size"):
48
+ self.config.vocab_size = self.get_input_embeddings().num_embeddings
49
+ self.vocab_size = self.config.vocab_size
50
+ if getattr(self.config, "tie_word_embeddings", True):
51
+ self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())
52
+
53
+ # Ensure 'lm_head.weight' exists when saving (avoids resume warnings)
54
+ def state_dict(self, *args, **kwargs):
55
+ sd = super().state_dict(*args, **kwargs)
56
+ if "lm_head.weight" not in sd and getattr(self.config, "tie_word_embeddings", True):
57
+ emb_key = f"{self.base_model_prefix}.embed_tokens.weight"
58
+ if emb_key in sd:
59
+ sd["lm_head.weight"] = sd[emb_key]
60
+ return sd
61
+
62
+ def forward(
63
+ self,
64
+ input_ids: Optional[torch.LongTensor] = None,
65
+ attention_mask: Optional[torch.Tensor] = None,
66
+ position_ids: Optional[torch.LongTensor] = None,
67
+ inputs_embeds: Optional[torch.FloatTensor] = None,
68
+ labels: Optional[torch.LongTensor] = None,
69
+ output_attentions: Optional[bool] = None,
70
+ output_hidden_states: Optional[bool] = None,
71
+ return_dict: Optional[bool] = True,
72
+ **kwargs,
73
+ ) -> Union[MaskedLMOutput, Tuple[torch.Tensor, ...]]:
74
+
75
+ outputs = self.encoder(
76
+ input_ids=input_ids,
77
+ attention_mask=attention_mask,
78
+ position_ids=position_ids,
79
+ inputs_embeds=inputs_embeds,
80
+ use_cache=False,
81
+ is_causal=False,
82
+ output_attentions=output_attentions,
83
+ output_hidden_states=output_hidden_states,
84
+ **kwargs,
85
+ )
86
+
87
+ hidden_states = outputs.last_hidden_state
88
+ logits = self.lm_head(hidden_states)
89
+
90
+ loss = None
91
+ if labels is not None:
92
+ loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
93
+ loss = loss_fct(logits.view(-1, self.vocab_size), labels.view(-1))
94
+
95
+ if not return_dict:
96
+ out = (logits, hidden_states)
97
+ if output_hidden_states:
98
+ out += (outputs.hidden_states,)
99
+ if output_attentions:
100
+ out += (outputs.attentions,)
101
+ if loss is not None:
102
+ out = (loss,) + out
103
+ return out
104
+
105
+ return MaskedLMOutput(
106
+ loss=loss,
107
+ logits=logits,
108
+ hidden_states=outputs.hidden_states,
109
+ attentions=outputs.attentions,
110
+ )
111
+
112
+
113
+ class Gemma3EncoderForSequenceClassification(Gemma3PreTrainedModel):
114
+ """Gemma3 Encoder with a sequence classification head (mean pooling + linear)."""
115
+ config_class = Gemma3TextConfig
116
+ base_model_prefix = "encoder"
117
+
118
+ def __init__(self, config: Gemma3TextConfig):
119
+ cfg = copy.deepcopy(config)
120
+ if hasattr(cfg, "use_bidirectional_attention"):
121
+ cfg.use_bidirectional_attention = True
122
+ cfg.use_cache = False
123
+ super().__init__(cfg)
124
+
125
+ self.num_labels = getattr(cfg, "num_labels", 2)
126
+ self.encoder = Gemma3TextModel(cfg)
127
+
128
+ classifier_dropout = getattr(cfg, "classifier_dropout", 0.0)
129
+ self.dropout = nn.Dropout(classifier_dropout)
130
+ self.classifier = nn.Linear(cfg.hidden_size, self.num_labels)
131
+
132
+ self.post_init()
133
+
134
+ def get_input_embeddings(self):
135
+ return self.encoder.embed_tokens
136
+
137
+ def set_input_embeddings(self, new_embeddings):
138
+ self.encoder.embed_tokens = new_embeddings
139
+
140
+ def forward(
141
+ self,
142
+ input_ids: Optional[torch.LongTensor] = None,
143
+ attention_mask: Optional[torch.Tensor] = None,
144
+ position_ids: Optional[torch.LongTensor] = None,
145
+ inputs_embeds: Optional[torch.FloatTensor] = None,
146
+ labels: Optional[torch.LongTensor] = None,
147
+ output_attentions: Optional[bool] = None,
148
+ output_hidden_states: Optional[bool] = None,
149
+ return_dict: Optional[bool] = True,
150
+ **kwargs,
151
+ ) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
152
+
153
+ outputs = self.encoder(
154
+ input_ids=input_ids,
155
+ attention_mask=attention_mask,
156
+ position_ids=position_ids,
157
+ inputs_embeds=inputs_embeds,
158
+ use_cache=False,
159
+ is_causal=False,
160
+ output_attentions=output_attentions,
161
+ output_hidden_states=output_hidden_states,
162
+ **kwargs,
163
+ )
164
+
165
+ hidden_states = outputs.last_hidden_state # (batch, seq_len, hidden)
166
+
167
+ # Mean pooling over non-padded tokens
168
+ if attention_mask is not None:
169
+ mask = attention_mask.unsqueeze(-1).float() # (batch, seq_len, 1)
170
+ pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
171
+ else:
172
+ pooled = hidden_states.mean(dim=1)
173
+
174
+ pooled = self.dropout(pooled)
175
+ logits = self.classifier(pooled)
176
+
177
+ loss = None
178
+ if labels is not None:
179
+ if self.config.problem_type is None:
180
+ if self.num_labels == 1:
181
+ self.config.problem_type = "regression"
182
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
183
+ self.config.problem_type = "single_label_classification"
184
+ else:
185
+ self.config.problem_type = "multi_label_classification"
186
+
187
+ if self.config.problem_type == "regression":
188
+ loss_fct = nn.MSELoss()
189
+ if self.num_labels == 1:
190
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
191
+ else:
192
+ loss = loss_fct(logits, labels)
193
+ elif self.config.problem_type == "single_label_classification":
194
+ loss_fct = nn.CrossEntropyLoss()
195
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
196
+ elif self.config.problem_type == "multi_label_classification":
197
+ loss_fct = nn.BCEWithLogitsLoss()
198
+ loss = loss_fct(logits, labels)
199
+
200
+ if not return_dict:
201
+ output = (logits,) + outputs[2:]
202
+ return ((loss,) + output) if loss is not None else output
203
+
204
+ return SequenceClassifierOutput(
205
+ loss=loss,
206
+ logits=logits,
207
+ hidden_states=outputs.hidden_states,
208
+ attentions=outputs.attentions,
209
+ )
210
+
211
+
212
+ class Gemma3EncoderForTokenClassification(Gemma3PreTrainedModel):
213
+ """Gemma3 Encoder with a token classification head for NER/POS tagging."""
214
+ config_class = Gemma3TextConfig
215
+ base_model_prefix = "encoder"
216
+
217
+ def __init__(self, config: Gemma3TextConfig):
218
+ cfg = copy.deepcopy(config)
219
+ if hasattr(cfg, "use_bidirectional_attention"):
220
+ cfg.use_bidirectional_attention = True
221
+ cfg.use_cache = False
222
+ super().__init__(cfg)
223
+
224
+ self.num_labels = getattr(cfg, "num_labels", 2)
225
+ self.encoder = Gemma3TextModel(cfg)
226
+
227
+ classifier_dropout = getattr(cfg, "classifier_dropout", 0.0)
228
+ self.dropout = nn.Dropout(classifier_dropout)
229
+ self.classifier = nn.Linear(cfg.hidden_size, self.num_labels)
230
+
231
+ self.post_init()
232
+
233
+ def get_input_embeddings(self):
234
+ return self.encoder.embed_tokens
235
+
236
+ def set_input_embeddings(self, new_embeddings):
237
+ self.encoder.embed_tokens = new_embeddings
238
+
239
+ def forward(
240
+ self,
241
+ input_ids: Optional[torch.LongTensor] = None,
242
+ attention_mask: Optional[torch.Tensor] = None,
243
+ position_ids: Optional[torch.LongTensor] = None,
244
+ inputs_embeds: Optional[torch.FloatTensor] = None,
245
+ labels: Optional[torch.LongTensor] = None,
246
+ output_attentions: Optional[bool] = None,
247
+ output_hidden_states: Optional[bool] = None,
248
+ return_dict: Optional[bool] = True,
249
+ **kwargs,
250
+ ) -> Union[TokenClassifierOutput, Tuple[torch.Tensor, ...]]:
251
+
252
+ outputs = self.encoder(
253
+ input_ids=input_ids,
254
+ attention_mask=attention_mask,
255
+ position_ids=position_ids,
256
+ inputs_embeds=inputs_embeds,
257
+ use_cache=False,
258
+ is_causal=False,
259
+ output_attentions=output_attentions,
260
+ output_hidden_states=output_hidden_states,
261
+ **kwargs,
262
+ )
263
+
264
+ hidden_states = outputs.last_hidden_state
265
+ hidden_states = self.dropout(hidden_states)
266
+ logits = self.classifier(hidden_states)
267
+
268
+ loss = None
269
+ if labels is not None:
270
+ loss_fct = nn.CrossEntropyLoss()
271
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
272
+
273
+ if not return_dict:
274
+ output = (logits,) + outputs[2:]
275
+ return ((loss,) + output) if loss is not None else output
276
+
277
+ return TokenClassifierOutput(
278
+ loss=loss,
279
+ logits=logits,
280
+ hidden_states=outputs.hidden_states,
281
+ attentions=outputs.attentions,
282
+ )
special_tokens_map.json ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "boi_token": "<start_of_image>",
3
+ "bos_token": {
4
+ "content": "<bos>",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false
9
+ },
10
+ "eoi_token": "<end_of_image>",
11
+ "eos_token": {
12
+ "content": "<eos>",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false
17
+ },
18
+ "image_token": "<image_soft_token>",
19
+ "mask_token": {
20
+ "content": "<mask>",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false
25
+ },
26
+ "pad_token": {
27
+ "content": "<pad>",
28
+ "lstrip": false,
29
+ "normalized": false,
30
+ "rstrip": false,
31
+ "single_word": false
32
+ },
33
+ "unk_token": {
34
+ "content": "<unk>",
35
+ "lstrip": false,
36
+ "normalized": false,
37
+ "rstrip": false,
38
+ "single_word": false
39
+ }
40
+ }
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:daab2354f8a74e70d70b4d1f804939b68a8c9624dd06cb7858e52dd8970e9726
3
+ size 33384567
tokenizer_config.json ADDED
The diff for this file is too large to render. See raw diff