sreddy109 commited on
Commit
12fd4a6
1 Parent(s): 8eb2c90

Add new SentenceTransformer model.

Browse files
1_Pooling/config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "word_embedding_dimension": 768,
3
+ "pooling_mode_cls_token": false,
4
+ "pooling_mode_mean_tokens": true,
5
+ "pooling_mode_max_tokens": false,
6
+ "pooling_mode_mean_sqrt_len_tokens": false,
7
+ "pooling_mode_weightedmean_tokens": false,
8
+ "pooling_mode_lasttoken": false
9
+ }
README.md ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: sentence-transformers
3
+ pipeline_tag: sentence-similarity
4
+ tags:
5
+ - sentence-transformers
6
+ - feature-extraction
7
+ - sentence-similarity
8
+
9
+ ---
10
+
11
+ # sreddy109/csbg-nomic-ai-upload-v2run5
12
+
13
+ This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 768 dimensional dense vector space and can be used for tasks like clustering or semantic search.
14
+
15
+ <!--- Describe your model here -->
16
+
17
+ ## Usage (Sentence-Transformers)
18
+
19
+ Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
20
+
21
+ ```
22
+ pip install -U sentence-transformers
23
+ ```
24
+
25
+ Then you can use the model like this:
26
+
27
+ ```python
28
+ from sentence_transformers import SentenceTransformer
29
+ sentences = ["This is an example sentence", "Each sentence is converted"]
30
+
31
+ model = SentenceTransformer('sreddy109/csbg-nomic-ai-upload-v2run5')
32
+ embeddings = model.encode(sentences)
33
+ print(embeddings)
34
+ ```
35
+
36
+
37
+
38
+ ## Evaluation Results
39
+
40
+ <!--- Describe how your model was evaluated -->
41
+
42
+ For an automated evaluation of this model, see the *Sentence Embeddings Benchmark*: [https://seb.sbert.net](https://seb.sbert.net?model_name=sreddy109/csbg-nomic-ai-upload-v2run5)
43
+
44
+
45
+ ## Training
46
+ The model was trained with the parameters:
47
+
48
+ **DataLoader**:
49
+
50
+ `torch.utils.data.dataloader.DataLoader` of length 1560 with parameters:
51
+ ```
52
+ {'batch_size': 8, 'sampler': 'torch.utils.data.sampler.RandomSampler', 'batch_sampler': 'torch.utils.data.sampler.BatchSampler'}
53
+ ```
54
+
55
+ **Loss**:
56
+
57
+ `sentence_transformers.losses.CachedMultipleNegativesRankingLoss.CachedMultipleNegativesRankingLoss` with parameters:
58
+ ```
59
+ {'scale': 20.0, 'similarity_fct': 'cos_sim'}
60
+ ```
61
+
62
+ Parameters of the fit()-Method:
63
+ ```
64
+ {
65
+ "epochs": 3,
66
+ "evaluation_steps": 78,
67
+ "evaluator": "sentence_transformers.evaluation.SequentialEvaluator.SequentialEvaluator",
68
+ "max_grad_norm": 1,
69
+ "optimizer_class": "<class 'torch.optim.adamw.AdamW'>",
70
+ "optimizer_params": {
71
+ "lr": 5e-07
72
+ },
73
+ "scheduler": "WarmupLinear",
74
+ "steps_per_epoch": null,
75
+ "warmup_steps": 1000,
76
+ "weight_decay": 0.01
77
+ }
78
+ ```
79
+
80
+
81
+ ## Full Model Architecture
82
+ ```
83
+ SentenceTransformer(
84
+ (0): Transformer({'max_seq_length': 1024, 'do_lower_case': False}) with Transformer model: NomicBertModel
85
+ (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False})
86
+ (2): Normalize()
87
+ )
88
+ ```
89
+
90
+ ## Citing & Authors
91
+
92
+ <!--- Describe where people can find more information -->
config.json ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "train_st_output/v2-run5/3900",
3
+ "activation_function": "swiglu",
4
+ "architectures": [
5
+ "NomicBertModel"
6
+ ],
7
+ "attn_pdrop": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_hf_nomic_bert.NomicBertConfig",
10
+ "AutoModel": "modeling_hf_nomic_bert.NomicBertModel",
11
+ "AutoModelForMaskedLM": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForPreTraining"
12
+ },
13
+ "bos_token_id": null,
14
+ "causal": false,
15
+ "dense_seq_output": true,
16
+ "embd_pdrop": 0.0,
17
+ "eos_token_id": null,
18
+ "fused_bias_fc": true,
19
+ "fused_dropout_add_ln": true,
20
+ "initializer_range": 0.02,
21
+ "layer_norm_epsilon": 1e-12,
22
+ "max_trained_positions": 2048,
23
+ "mlp_fc1_bias": false,
24
+ "mlp_fc2_bias": false,
25
+ "model_type": "nomic_bert",
26
+ "n_embd": 768,
27
+ "n_head": 12,
28
+ "n_inner": 3072,
29
+ "n_layer": 12,
30
+ "n_positions": 8192,
31
+ "pad_vocab_size_multiple": 64,
32
+ "parallel_block": false,
33
+ "parallel_block_tied_norm": false,
34
+ "prenorm": false,
35
+ "qkv_proj_bias": false,
36
+ "reorder_and_upcast_attn": false,
37
+ "resid_pdrop": 0.0,
38
+ "rotary_emb_base": 1000,
39
+ "rotary_emb_fraction": 1.0,
40
+ "rotary_emb_interleaved": false,
41
+ "rotary_emb_scale_base": null,
42
+ "rotary_scaling_factor": 2,
43
+ "scale_attn_by_inverse_layer_idx": false,
44
+ "scale_attn_weights": true,
45
+ "summary_activation": null,
46
+ "summary_first_dropout": 0.0,
47
+ "summary_proj_to_labels": true,
48
+ "summary_type": "cls_index",
49
+ "summary_use_proj": true,
50
+ "torch_dtype": "float32",
51
+ "transformers_version": "4.37.2",
52
+ "type_vocab_size": 2,
53
+ "use_cache": true,
54
+ "use_flash_attn": true,
55
+ "use_rms_norm": false,
56
+ "use_xentropy": true,
57
+ "vocab_size": 30528
58
+ }
config_sentence_transformers.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "__version__": {
3
+ "sentence_transformers": "2.4.0.dev0",
4
+ "transformers": "4.37.2",
5
+ "pytorch": "2.1.0+cu121"
6
+ }
7
+ }
configuration_hf_nomic_bert.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import GPT2Config
2
+
3
+
4
+ class NomicBertConfig(GPT2Config):
5
+ model_type = "nomic_bert"
6
+
7
+ def __init__(
8
+ self,
9
+ prenorm=False,
10
+ parallel_block=False,
11
+ parallel_block_tied_norm=False,
12
+ rotary_emb_fraction=0.0,
13
+ fused_dropout_add_ln=False,
14
+ fused_bias_fc=False,
15
+ use_flash_attn=False,
16
+ use_xentropy=False,
17
+ qkv_proj_bias=True,
18
+ rotary_emb_base=1000,
19
+ rotary_emb_scale_base=None,
20
+ rotary_emb_interleaved=False,
21
+ mlp_fc1_bias=True,
22
+ mlp_fc2_bias=True,
23
+ use_rms_norm=False,
24
+ causal=False,
25
+ type_vocab_size=2,
26
+ dense_seq_output=True,
27
+ pad_vocab_size_multiple=1,
28
+ tie_word_embeddings=True,
29
+ rotary_scaling_factor=1.0,
30
+ max_trained_positions=2048,
31
+ **kwargs,
32
+ ):
33
+ self.prenorm = prenorm
34
+ self.parallel_block = parallel_block
35
+ self.parallel_block_tied_norm = parallel_block_tied_norm
36
+ self.rotary_emb_fraction = rotary_emb_fraction
37
+ self.tie_word_embeddings = tie_word_embeddings
38
+ self.fused_dropout_add_ln = fused_dropout_add_ln
39
+ self.fused_bias_fc = fused_bias_fc
40
+ self.use_flash_attn = use_flash_attn
41
+ self.use_xentropy = use_xentropy
42
+ self.qkv_proj_bias = qkv_proj_bias
43
+ self.rotary_emb_base = rotary_emb_base
44
+ self.rotary_emb_scale_base = rotary_emb_scale_base
45
+ self.rotary_emb_interleaved = rotary_emb_interleaved
46
+ self.mlp_fc1_bias = mlp_fc1_bias
47
+ self.mlp_fc2_bias = mlp_fc2_bias
48
+ self.use_rms_norm = use_rms_norm
49
+ self.causal = causal
50
+ self.type_vocab_size = type_vocab_size
51
+ self.dense_seq_output = dense_seq_output
52
+ self.pad_vocab_size_multiple = pad_vocab_size_multiple
53
+ self.rotary_scaling_factor = rotary_scaling_factor
54
+ self.max_trained_positions = max_trained_positions
55
+
56
+ super().__init__(**kwargs)
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e5f6bcf5439a8393c75043c6ef244f83bc415915b946b773a6668b37af4fc842
3
+ size 546938168
modeling_hf_nomic_bert.py ADDED
@@ -0,0 +1,1221 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2022, Tri Dao.
2
+ # This BERT implementation is based on our MLPerf 2.0 and MLPerf 2.1 BERT implementation.
3
+ # https://github.com/mlcommons/training_results_v2.0/blob/main/HazyResearch/benchmarks/bert/implementations/pytorch/modeling.py
4
+ # https://github.com/mlcommons/training_results_v2.1/blob/main/Azure-HazyResearch/benchmarks/bert/implementations/ND96amsr_A100_v4/modeling.py
5
+
6
+ import logging
7
+
8
+ # Inspired by https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py
9
+ import os
10
+ import re
11
+ from collections import OrderedDict
12
+ from functools import partial
13
+ from typing import List, Optional, Tuple, Union
14
+
15
+ import torch
16
+ import torch.nn as nn
17
+ import torch.nn.functional as F
18
+ from einops import rearrange, repeat
19
+ from safetensors.torch import load_file as safe_load_file
20
+ from transformers import GPT2Config, PreTrainedModel
21
+ from transformers.models.bert.modeling_bert import (
22
+ BaseModelOutputWithPoolingAndCrossAttentions,
23
+ MaskedLMOutput,
24
+ SequenceClassifierOutput,
25
+ )
26
+ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME
27
+ from transformers.utils.hub import cached_file, get_checkpoint_shard_files
28
+
29
+ from .configuration_hf_nomic_bert import NomicBertConfig
30
+
31
+ logger = logging.getLogger(__name__)
32
+
33
+
34
+ # adapted from flash attention, added safe serialization option for hf models
35
+ def state_dict_from_pretrained(model_name, safe_serialization=False, device=None, dtype=None):
36
+ # If not fp32, then we don't want to load directly to the GPU
37
+ mapped_device = "cpu" if dtype not in [torch.float32, None] else device
38
+ is_sharded = False
39
+ load_safe = False
40
+ resolved_archive_file = None
41
+
42
+ weights_path = os.path.join(model_name, WEIGHTS_NAME)
43
+ weights_index_path = os.path.join(model_name, WEIGHTS_INDEX_NAME)
44
+ safe_weights_path = os.path.join(model_name, SAFE_WEIGHTS_NAME)
45
+ safe_weights_index_path = os.path.join(model_name, SAFE_WEIGHTS_INDEX_NAME)
46
+
47
+ if os.path.isfile(weights_path):
48
+ resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
49
+ elif os.path.isfile(weights_index_path):
50
+ resolved_archive_file = cached_file(model_name, WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False)
51
+ is_sharded = True
52
+ elif os.path.isfile(safe_weights_path):
53
+ resolved_archive_file = cached_file(model_name, SAFE_WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
54
+ load_safe = True
55
+ elif os.path.isfile(safe_weights_index_path):
56
+ resolved_archive_file = cached_file(
57
+ model_name, SAFE_WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False
58
+ )
59
+ is_sharded = True
60
+ load_safe = True
61
+ else: # Try loading from HF hub instead of from local files
62
+ weight_name = WEIGHTS_NAME if not safe_serialization else SAFE_WEIGHTS_NAME
63
+ resolved_archive_file = cached_file(model_name, weight_name, _raise_exceptions_for_missing_entries=False)
64
+ if resolved_archive_file is None:
65
+ weight_index = WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_WEIGHTS_INDEX_NAME
66
+ resolved_archive_file = cached_file(model_name, weight_index, _raise_exceptions_for_missing_entries=False)
67
+ if resolved_archive_file is not None:
68
+ is_sharded = True
69
+
70
+ load_safe = safe_serialization
71
+
72
+ if resolved_archive_file is None:
73
+ raise EnvironmentError(f"Model name {model_name} was not found.")
74
+
75
+ if load_safe:
76
+ loader = partial(safe_load_file, device=mapped_device)
77
+ else:
78
+ loader = partial(torch.load, map_location=mapped_device)
79
+
80
+ if is_sharded:
81
+ # resolved_archive_file becomes a list of files that point to the different
82
+ # checkpoint shards in this case.
83
+ resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(model_name, resolved_archive_file)
84
+ state_dict = {}
85
+ for sharded_file in resolved_archive_file:
86
+ state_dict.update(loader(sharded_file))
87
+ else:
88
+ state_dict = loader(resolved_archive_file)
89
+ # Convert dtype before moving to GPU to save memory
90
+ if dtype is not None:
91
+ state_dict = {k: v.to(dtype=dtype) for k, v in state_dict.items()}
92
+ state_dict = {k: v.to(device=device) for k, v in state_dict.items()}
93
+ return state_dict
94
+
95
+
96
+ def filter_shapes(state_dict, model):
97
+ """
98
+ Filters the state dict to match the current model shape.
99
+ """
100
+ filtered_state_dict = {}
101
+ for key, value in state_dict.items():
102
+ if key in model.state_dict():
103
+ if value.shape == model.state_dict()[key].shape:
104
+ filtered_state_dict[key] = value
105
+ return filtered_state_dict
106
+
107
+
108
+ def remap_bert_state_dict(state_dict, config, remove_bert=False, remove_cls_weights=False, add_pooling_layer=False):
109
+ """
110
+ Map the state_dict of a Huggingface BERT model to be flash_attn compatible.
111
+ """
112
+
113
+ def add_bert_prefix(key):
114
+ # prepend bert. to the key
115
+ if key.startswith("bert.") or key.startswith("cls."):
116
+ return key
117
+ return f"bert.{key}"
118
+
119
+ state_dict = OrderedDict((add_bert_prefix(k), v) for k, v in state_dict.items())
120
+
121
+ # LayerNorm
122
+ def key_mapping_ln_gamma_beta(key):
123
+ key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key)
124
+ key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key)
125
+ return key
126
+
127
+ state_dict = OrderedDict((key_mapping_ln_gamma_beta(k), v) for k, v in state_dict.items())
128
+
129
+ # Layers
130
+ def key_mapping_layers(key):
131
+ return re.sub(r"^bert.encoder.layer\.", "bert.encoder.layers.", key)
132
+
133
+ state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items())
134
+
135
+ # LayerNorm
136
+ def key_mapping_ln(key):
137
+ key = re.sub(r"^bert.embeddings.LayerNorm.", "bert.emb_ln.", key)
138
+ key = re.sub(
139
+ r"^bert.encoder.layers.(\d+).attention.output.LayerNorm.(weight|bias)",
140
+ r"bert.encoder.layers.\1.norm1.\2",
141
+ key,
142
+ )
143
+ key = re.sub(
144
+ r"^bert.encoder.layers.(\d+).output.LayerNorm.(weight|bias)",
145
+ r"bert.encoder.layers.\1.norm2.\2",
146
+ key,
147
+ )
148
+ key = re.sub(
149
+ r"^cls.predictions.transform.LayerNorm.(weight|bias)",
150
+ r"cls.predictions.transform.layer_norm.\1",
151
+ key,
152
+ )
153
+ return key
154
+
155
+ state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
156
+
157
+ # MLP
158
+ def key_mapping_mlp(key):
159
+ key = re.sub(
160
+ r"^bert.encoder.layers.(\d+).intermediate.dense.(weight|bias)",
161
+ r"bert.encoder.layers.\1.mlp.fc1.\2",
162
+ key,
163
+ )
164
+ key = re.sub(
165
+ r"^bert.encoder.layers.(\d+).output.dense.(weight|bias)",
166
+ r"bert.encoder.layers.\1.mlp.fc2.\2",
167
+ key,
168
+ )
169
+ return key
170
+
171
+ state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())
172
+
173
+ # Attention
174
+ last_layer_subset = getattr(config, "last_layer_subset", False)
175
+ for d in range(config.num_hidden_layers):
176
+ if f"bert.encoder.layers.{d}.attention.self.query.weight" not in state_dict:
177
+ continue
178
+ Wq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.weight")
179
+ Wk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.weight")
180
+ Wv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.weight")
181
+ bq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.bias")
182
+ bk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.bias")
183
+ bv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.bias")
184
+ if not (last_layer_subset and d == config.num_hidden_layers - 1):
185
+ state_dict[f"bert.encoder.layers.{d}.attn.Wqkv.weight"] = torch.cat([Wq, Wk, Wv], dim=0)
186
+ state_dict[f"bert.encoder.layers.{d}.attn.Wqkv.bias"] = torch.cat([bq, bk, bv], dim=0)
187
+ else:
188
+ state_dict[f"bert.encoder.layers.{d}.attn.Wq.weight"] = Wq
189
+ state_dict[f"bert.encoder.layers.{d}.attn.Wkv.weight"] = torch.cat([Wk, Wv], dim=0)
190
+ state_dict[f"bert.encoder.layers.{d}.attn.Wq.bias"] = bq
191
+ state_dict[f"bert.encoder.layers.{d}.attn.Wkv.bias"] = torch.cat([bk, bv], dim=0)
192
+
193
+ def key_mapping_attn(key):
194
+ return re.sub(
195
+ r"^bert.encoder.layers.(\d+).attention.output.dense.(weight|bias)",
196
+ r"bert.encoder.layers.\1.attn.out_proj.\2",
197
+ key,
198
+ )
199
+
200
+ state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())
201
+
202
+ def key_mapping_decoder_bias(key):
203
+ return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key)
204
+
205
+ # remove nsp weights, we don't use
206
+ state_dict.pop("cls.seq_relationship.weight", None)
207
+ state_dict.pop("cls.seq_relationship.bias", None)
208
+ state_dict.pop("bert.embeddings.position_ids", None)
209
+
210
+ state_dict = OrderedDict((key_mapping_decoder_bias(k), v) for k, v in state_dict.items())
211
+
212
+ if remove_cls_weights:
213
+ cls_weights = [
214
+ "cls.predictions.decoder.bias",
215
+ "cls.predictions.transform.dense.weight",
216
+ "cls.predictions.transform.dense.bias",
217
+ "cls.predictions.transform.layer_norm.weight",
218
+ "cls.predictions.transform.layer_norm.bias",
219
+ "cls.predictions.decoder.weight",
220
+ ]
221
+ for weight in cls_weights:
222
+ state_dict.pop(weight, None)
223
+
224
+ # Word embedding
225
+ pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
226
+ if pad_vocab_size_multiple > 1:
227
+ word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
228
+ state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(
229
+ word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])
230
+ )
231
+ if not remove_cls_weights:
232
+ decoder_weight = state_dict["cls.predictions.decoder.weight"]
233
+ state_dict["cls.predictions.decoder.weight"] = F.pad(
234
+ decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])
235
+ )
236
+ # If the vocab was padded, we want to set the decoder bias for those padded indices to be
237
+ # strongly negative (i.e. the decoder shouldn't predict those indices).
238
+ # TD [2022-05-09]: I don't think it affects the MLPerf training.
239
+ if "cls.predictions.decoder.bias" in state_dict:
240
+ decoder_bias = state_dict["cls.predictions.decoder.bias"]
241
+ state_dict["cls.predictions.decoder.bias"] = F.pad(
242
+ decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
243
+ )
244
+
245
+ if add_pooling_layer is False:
246
+ pooler_weights = [
247
+ "bert.pooler.dense.weight",
248
+ "bert.pooler.dense.bias",
249
+ ]
250
+ for key in pooler_weights:
251
+ state_dict.pop(key, None)
252
+
253
+ if remove_bert:
254
+
255
+ def remove_bert_prefix(key):
256
+ key = re.sub(r"^bert.", "", key)
257
+ return key
258
+
259
+ state_dict = OrderedDict((remove_bert_prefix(k), v) for k, v in state_dict.items())
260
+
261
+ return state_dict
262
+
263
+
264
+ class NomicBertPreTrainedModel(PreTrainedModel):
265
+ """An abstract class to handle weights initialization and
266
+ a simple interface for dowloading and loading pretrained models.
267
+ """
268
+
269
+ config_class = NomicBertConfig
270
+ base_model_prefix = "model"
271
+ supports_gradient_checkpointing = True
272
+ _no_split_modules = ["Block"]
273
+ _skip_keys_device_placement = "past_key_values"
274
+
275
+ def __init__(self, config, *inputs, **kwargs):
276
+ super().__init__(config)
277
+ if not isinstance(config, GPT2Config):
278
+ raise ValueError(
279
+ "Parameter config in `{}(config)` should be an instance of class `GPT2Config`. "
280
+ "To create a model from a Google pretrained model use "
281
+ "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
282
+ self.__class__.__name__, self.__class__.__name__
283
+ )
284
+ )
285
+ self.config = config
286
+
287
+ @classmethod
288
+ def from_pretrained(cls, model_name, config=None, *inputs, **kwargs):
289
+ """
290
+ Instantiate a NomicBertPreTrainedModel from a pre-trained model file or a pytorch state dict.
291
+ Download and cache the pre-trained model file if needed.
292
+
293
+ Params:
294
+ pretrained_model_name_or_path: either:
295
+ - a path or url to a pretrained model archive containing:
296
+ . `bert_config.json` a configuration file for the model
297
+ . `pytorch_model.bin` a PyTorch dump of a NomicBertForPretraining instance
298
+ - a path or url to a pretrained model archive containing:
299
+ . `bert_config.json` a configuration file for the model
300
+ . `model.chkpt` a TensorFlow checkpoint
301
+ *inputs, **kwargs: additional input for the specific NomicBert class
302
+ (ex: num_labels for NomicBertForSequenceClassification)
303
+ """
304
+ # Instantiate model.
305
+ if config is None:
306
+ config = cls.config_class.from_pretrained(model_name)
307
+ remove_cls = cls != NomicBertForPreTraining
308
+ remove_bert_prefix = cls != NomicBertForPreTraining
309
+ ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
310
+ num_labels = kwargs.pop("num_labels", None)
311
+ rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
312
+ if rotary_scaling_factor:
313
+ config.rotary_scaling_factor = rotary_scaling_factor
314
+
315
+ if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
316
+ config.n_positions = 2048
317
+ if num_labels:
318
+ config.num_labels = num_labels
319
+
320
+ if "add_pooling_layer" in kwargs:
321
+ model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
322
+ else:
323
+ if cls == NomicBertModel:
324
+ model = cls(config, *inputs, add_pooling_layer=False)
325
+ else:
326
+ model = cls(config, *inputs)
327
+ # TODO: fix this
328
+ # Assuming we know what we're doing when loading from disk
329
+ # Prob a bad assumption but i'm tired and want to train this asap
330
+ if os.path.exists(model_name):
331
+ model_path = f"{model_name}/pytorch_model.bin"
332
+ if os.path.exists(model_path):
333
+ state_dict = torch.load(f"{model_name}/pytorch_model.bin")
334
+ else:
335
+ model_path = f"{model_name}/model.safetensors"
336
+ if not os.path.exists(model_path):
337
+ raise ValueError(f"Model path {model_path} not found")
338
+ state_dict = safe_load_file(model_path)
339
+
340
+ if ignore_mismatched_shapes:
341
+ state_dict = filter_shapes(state_dict, model)
342
+ load_return = model.load_state_dict(state_dict, strict=False)
343
+ else:
344
+ # TODO: can probably check config class and see if we need to remap from a bert model
345
+ state_dict = state_dict_from_pretrained(model_name, safe_serialization=kwargs.get("safe_serialization", False))
346
+ state_dict = remap_bert_state_dict(
347
+ state_dict,
348
+ config,
349
+ remove_bert=remove_bert_prefix,
350
+ remove_cls_weights=remove_cls,
351
+ add_pooling_layer=getattr(config, "add_pooling_layer", False),
352
+ )
353
+ if ignore_mismatched_shapes:
354
+ state_dict = filter_shapes(state_dict, model)
355
+
356
+ load_return = model.load_state_dict(state_dict, strict=True)
357
+ logger.warning(load_return)
358
+ return model
359
+
360
+ def _set_gradient_checkpointing(self, module, value=False):
361
+ if isinstance(module, NomicBertEncoder):
362
+ module.gradient_checkpointing = value
363
+
364
+
365
+ # https://github.com/huggingface/transformers/blob/7032e0203262ebb2ebf55da8d2e01f873973e835/src/transformers/models/bert/modeling_bert.py#L748
366
+ def _init_weights(module, initializer_range=0.02):
367
+ if isinstance(module, nn.Linear):
368
+ nn.init.normal_(module.weight, std=initializer_range)
369
+ if module.bias is not None:
370
+ nn.init.zeros_(module.bias)
371
+ elif isinstance(module, nn.Embedding):
372
+ nn.init.normal_(module.weight, std=initializer_range)
373
+ if module.padding_idx is not None:
374
+ nn.init.zeros_(module.weight[module.padding_idx])
375
+
376
+
377
+ class NomicBertEmbeddings(nn.Module):
378
+ def __init__(self, config):
379
+ """
380
+ If max_position_embeddings <= 0, there's no position embeddings
381
+ If type_vocab_size <= 0, there's no token type embeddings
382
+ """
383
+ super().__init__()
384
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
385
+ self.max_position_embeddings = config.max_position_embeddings if config.rotary_emb_fraction <= 0 else 0
386
+ self.type_vocab_size = config.type_vocab_size
387
+ if self.max_position_embeddings > 0 and config.rotary_emb_fraction <= 0:
388
+ self.position_embeddings = nn.Embedding(
389
+ config.max_position_embeddings,
390
+ config.hidden_size,
391
+ )
392
+ if self.type_vocab_size > 0:
393
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
394
+
395
+ def forward(self, input_ids, position_ids=None, token_type_ids=None):
396
+ """
397
+ input_ids: (batch, seqlen)
398
+ position_ids: (batch, seqlen)
399
+ token_type_ids: (batch, seqlen)
400
+ """
401
+ batch_size, seqlen = input_ids.shape
402
+ embeddings = self.word_embeddings(input_ids)
403
+
404
+ if self.type_vocab_size > 0:
405
+ if token_type_ids is None:
406
+ token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=input_ids.device)
407
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
408
+ embeddings = embeddings + token_type_embeddings
409
+
410
+ if self.max_position_embeddings > 0:
411
+ if position_ids is None:
412
+ position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
413
+ position_embeddings = self.position_embeddings(position_ids)
414
+ embeddings = embeddings + position_embeddings
415
+ return embeddings
416
+
417
+
418
+ class NomicBertMLP(nn.Module):
419
+ def __init__(
420
+ self,
421
+ in_features,
422
+ hidden_features=None,
423
+ out_features=None,
424
+ activation=F.gelu,
425
+ bias1=True,
426
+ bias2=True,
427
+ return_residual=False,
428
+ fused_bias_fc=False,
429
+ ):
430
+ super().__init__()
431
+ out_features = out_features if out_features is not None else in_features
432
+ hidden_features = hidden_features if hidden_features is not None else in_features * 4
433
+ self.return_residual = return_residual
434
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1)
435
+ approximate = "tanh" if activation in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] else "none"
436
+ self.activation = nn.GELU(approximate=approximate) if activation == "gelu" else activation
437
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)
438
+
439
+ def forward(self, x):
440
+ y = self.fc1(x)
441
+ y = self.activation(y)
442
+ y = self.fc2(y)
443
+ return y if not self.return_residual else (y, x)
444
+
445
+
446
+ class NomciBertGatedMLP(nn.Module):
447
+ def __init__(
448
+ self,
449
+ in_features,
450
+ hidden_features=None,
451
+ out_features=None,
452
+ activation=F.sigmoid,
453
+ bias1=True,
454
+ bias2=True,
455
+ multiple_of=256,
456
+ return_residual=False,
457
+ fused_bias_fc=True,
458
+ device=None,
459
+ dtype=None,
460
+ ):
461
+ super().__init__()
462
+ out_features = out_features if out_features is not None else in_features
463
+ hidden_features = hidden_features if hidden_features is not None else int(8 * in_features / 3)
464
+ hidden_features = (hidden_features + multiple_of - 1) // multiple_of * multiple_of
465
+ self.return_residual = return_residual
466
+
467
+ self.fc11 = nn.Linear(in_features, hidden_features, bias=bias1)
468
+ self.fc12 = nn.Linear(in_features, hidden_features, bias=bias1)
469
+ self.activation = activation
470
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)
471
+
472
+ def forward(self, x):
473
+ y = self.fc11(x)
474
+ gate = self.fc12(x)
475
+ if self.activation == F.sigmoid: # Special case for GLU
476
+ y = F.glu(torch.cat([y, gate], dim=-1), dim=-1)
477
+ else:
478
+ y = y * self.activation(gate)
479
+ y = self.fc2(y)
480
+ return y if not self.return_residual else (y, x)
481
+
482
+
483
+ def rotate_half(x, interleaved=False):
484
+ if not interleaved:
485
+ x1, x2 = x.chunk(2, dim=-1)
486
+ return torch.cat((-x2, x1), dim=-1)
487
+ else:
488
+ x1, x2 = x[..., ::2], x[..., 1::2]
489
+ return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
490
+
491
+
492
+ def apply_rotary_emb(x, cos, sin, offset=0, interleaved=False):
493
+ """
494
+ x: (batch_size, seqlen, nheads, headdim)
495
+ cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
496
+ """
497
+ ro_dim = cos.shape[-1] * 2
498
+ assert ro_dim <= x.shape[-1]
499
+ cos, sin = (
500
+ cos[offset : offset + x.shape[1]],
501
+ sin[offset : offset + x.shape[1]],
502
+ )
503
+ cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
504
+ sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
505
+ return torch.cat(
506
+ [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]],
507
+ dim=-1,
508
+ )
509
+
510
+
511
+ class NomicBertRotaryEmbedding(nn.Module):
512
+ def __init__(
513
+ self,
514
+ dim: int,
515
+ base=10000.0,
516
+ interleaved=False,
517
+ scale_base=None,
518
+ pos_idx_in_fp32=True,
519
+ device=None,
520
+ ):
521
+ """
522
+ interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
523
+ of 1st half and 2nd half (GPT-NeoX style).
524
+ pos_idx_in_fp32: if True, the position indices [0.0, ..., seqlen - 1] are in fp32,
525
+ otherwise they might be in lower precision.
526
+ This option was added because previously (before 2023-07-02), when we construct
527
+ the position indices, we use the dtype of self.inv_freq. In most cases this would
528
+ be fp32, but if the model is trained in pure bf16 (not mixed precision), then
529
+ self.inv_freq would be bf16, and the position indices are also in bf16.
530
+ Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the
531
+ embeddings for some positions will coincide.
532
+ To maintain compatibility with models previously trained in pure bf16,
533
+ we add this option.
534
+ """
535
+ super().__init__()
536
+ self.dim = dim
537
+ self.base = float(base)
538
+ self.pos_idx_in_fp32 = pos_idx_in_fp32
539
+ # Generate and save the inverse frequency buffer (non trainable)
540
+ inv_freq = self._compute_inv_freq(device)
541
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
542
+ self.interleaved = interleaved
543
+ self.scale_base = scale_base
544
+ scale = (
545
+ (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
546
+ if scale_base is not None
547
+ else None
548
+ )
549
+ self.register_buffer("scale", scale, persistent=False)
550
+
551
+ self._seq_len_cached = 0
552
+ self._cos_cached = None
553
+ self._sin_cached = None
554
+ self._cos_k_cached = None
555
+ self._sin_k_cached = None
556
+
557
+ def _compute_inv_freq(self, device=None):
558
+ return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
559
+
560
+ def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
561
+ # Reset the tables if the sequence length has changed,
562
+ # if we're on a new device (possibly due to tracing for instance),
563
+ # or if we're switching from inference mode to training
564
+ if (
565
+ seqlen > self._seq_len_cached
566
+ or self._cos_cached is None
567
+ or self._cos_cached.device != device
568
+ or self._cos_cached.dtype != dtype
569
+ or (self.training and self._cos_cached.is_inference())
570
+ ):
571
+ self._seq_len_cached = seqlen
572
+ # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
573
+ # And the output of arange can be quite large, so bf16 would lose a lot of precision.
574
+ # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
575
+ if self.pos_idx_in_fp32:
576
+ t = torch.arange(seqlen, device=device, dtype=torch.float32)
577
+ # We want fp32 here as well since inv_freq will be multiplied with t, and the output
578
+ # will be large. Having it in bf16 will lose a lot of precision and cause the
579
+ # cos & sin output to change significantly.
580
+ # We want to recompute self.inv_freq if it was not loaded in fp32
581
+ if self.inv_freq.dtype != torch.float32:
582
+ inv_freq = self._compute_inv_freq(device=device)
583
+ else:
584
+ inv_freq = self.inv_freq
585
+ else:
586
+ t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
587
+ inv_freq = self.inv_freq
588
+ # Don't do einsum, it converts fp32 to fp16 under AMP
589
+ # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
590
+ freqs = torch.outer(t, inv_freq)
591
+ self._cos_cached = torch.cos(freqs).to(dtype)
592
+ self._sin_cached = torch.sin(freqs).to(dtype)
593
+
594
+ def forward(
595
+ self,
596
+ qkv: torch.Tensor,
597
+ kv: Optional[torch.Tensor] = None,
598
+ seqlen_offset: Union[int, torch.Tensor] = 0,
599
+ max_seqlen: Optional[int] = None,
600
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
601
+ """
602
+ qkv: (batch, seqlen, 3, nheads, headdim) if kv is none,
603
+ else it's just q of shape (batch, seqlen, nheads, headdim)
604
+ kv: (batch, seqlen, 2, nheads, headdim)
605
+ seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount.
606
+ Most commonly used in inference when we have KV cache.
607
+ If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one
608
+ should pass in max_seqlen, which will update the cos / sin cache up to that length.
609
+ Apply rotary embedding *inplace* to qkv and / or kv.
610
+ """
611
+ seqlen = qkv.shape[1]
612
+ if seqlen > self._seq_len_cached:
613
+ self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
614
+ elif max_seqlen is not None:
615
+ self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
616
+ elif isinstance(seqlen_offset, int):
617
+ self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
618
+
619
+ q_rot = apply_rotary_emb(qkv[:, :, 0], self._cos_cached, self._sin_cached, seqlen_offset, self.interleaved)
620
+ k_rot = apply_rotary_emb(qkv[:, :, 1], self._cos_cached, self._sin_cached, seqlen_offset, self.interleaved)
621
+ return torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
622
+
623
+
624
+ class NomicBertDynamicNTKRotaryEmbedding(NomicBertRotaryEmbedding):
625
+ def __init__(self, rotary_scaling_factor, max_position_embeddings, **kwargs):
626
+ super().__init__(**kwargs)
627
+ self.rotary_scaling_factor = rotary_scaling_factor
628
+ self.max_position_embeddings = max_position_embeddings
629
+
630
+ def _compute_inv_freq(self, base=None, device=None):
631
+ if base is None:
632
+ base = self.base
633
+ return 1.0 / (base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
634
+
635
+ def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
636
+ # Reset the tables if the sequence length has changed,
637
+ # if we're on a new device (possibly due to tracing for instance),
638
+ # or if we're switching from inference mode to training
639
+ if seqlen > self.max_position_embeddings:
640
+ base = self.base * (
641
+ (self.rotary_scaling_factor * seqlen / self.max_position_embeddings) - (self.rotary_scaling_factor - 1)
642
+ ) ** (self.dim / (self.dim - 2))
643
+ inv_freq = self._compute_inv_freq(base=base, device=device)
644
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
645
+
646
+ if (
647
+ seqlen > self._seq_len_cached
648
+ or self._cos_cached is None
649
+ or self._cos_cached.device != device
650
+ or self._cos_cached.dtype != dtype
651
+ or (self.training and self._cos_cached.is_inference())
652
+ ):
653
+ self._seq_len_cached = seqlen
654
+ # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
655
+ # And the output of arange can be quite large, so bf16 would lose a lot of precision.
656
+ # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
657
+ if self.pos_idx_in_fp32:
658
+ t = torch.arange(seqlen, device=device, dtype=torch.float32)
659
+ # We want fp32 here as well since inv_freq will be multiplied with t, and the output
660
+ # will be large. Having it in bf16 will lose a lot of precision and cause the
661
+ # cos & sin output to change significantly.
662
+ # We want to recompute self.inv_freq if it was not loaded in fp32
663
+ if self.inv_freq.dtype != torch.float32:
664
+ if seqlen > self.max_position_embeddings:
665
+ base = self.base * (
666
+ (self.scaling_factor * seqlen / self.max_position_embeddings) - (self.scaling_factor - 1)
667
+ ) ** (self.dim / (self.dim - 2))
668
+ else:
669
+ base = self.base
670
+ inv_freq = self._compute_inv_freq(device=device, base=base)
671
+ else:
672
+ inv_freq = self.inv_freq
673
+ else:
674
+ t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
675
+ inv_freq = self.inv_freq
676
+ # Don't do einsum, it converts fp32 to fp16 under AMP
677
+ # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
678
+ freqs = torch.outer(t, inv_freq)
679
+ if self.scale is None:
680
+ self._cos_cached = torch.cos(freqs).to(dtype)
681
+ self._sin_cached = torch.sin(freqs).to(dtype)
682
+ else:
683
+ power = (
684
+ torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
685
+ ) / self.scale_base
686
+ scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
687
+ # We want the multiplication by scale to happen in fp32
688
+ self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
689
+ self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
690
+ self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
691
+ self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
692
+
693
+
694
+ class NomicBertAttention(nn.Module):
695
+ """Multi-head self-attention and cross-attention"""
696
+
697
+ def __init__(
698
+ self,
699
+ config,
700
+ ) -> None:
701
+ """
702
+ num_heads_kv: can be used to toggle MQA / GQA. If None, use num_heads.
703
+ return_residual: whether to return the input x along with the output. This is for
704
+ performance reason: for post-norm architecture, returning the input allows us
705
+ to fuse the backward of nn.Linear with the residual connection.
706
+ """
707
+ super().__init__()
708
+ self.embed_dim = config.n_embd
709
+ self.use_flash_attn = config.use_flash_attn
710
+ self.fused_bias_fc = config.fused_bias_fc
711
+
712
+ self.num_heads = config.n_head
713
+ self.num_heads_kv = config.num_heads_kv if getattr(config, "num_heads_kv", None) is not None else self.num_heads
714
+ assert self.embed_dim % self.num_heads == 0, "embed_dim must be divisible by num_heads"
715
+ self.head_dim = self.embed_dim // self.num_heads
716
+ # we don't really support mqa / gqa for now
717
+ qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_heads_kv)
718
+
719
+ self.register_buffer(
720
+ "norm_factor",
721
+ torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
722
+ persistent=False,
723
+ )
724
+
725
+ self.rotary_emb_dim = self.head_dim * config.rotary_emb_fraction
726
+ if self.rotary_emb_dim > 0:
727
+ if config.rotary_scaling_factor:
728
+ self.rotary_emb = NomicBertDynamicNTKRotaryEmbedding(
729
+ dim=self.rotary_emb_dim,
730
+ base=config.rotary_emb_base,
731
+ scale_base=config.rotary_emb_scale_base,
732
+ interleaved=config.rotary_emb_interleaved,
733
+ rotary_scaling_factor=config.rotary_scaling_factor,
734
+ max_position_embeddings=config.max_trained_positions,
735
+ )
736
+ else:
737
+ self.rotary_emb = NomicBertRotaryEmbedding(
738
+ dim=self.rotary_emb_dim,
739
+ base=config.rotary_emb_base,
740
+ scale_base=config.rotary_emb_scale_base,
741
+ interleaved=config.rotary_emb_interleaved,
742
+ )
743
+ # bug in xformers: https://github.com/facebookresearch/xformers/issues/841
744
+ # uses the head dimension instead of the sequence dimension
745
+ self.rotary_head_dim = getattr(config, "rotary_head_dim", False)
746
+
747
+ self.Wqkv = nn.Linear(self.embed_dim, qkv_dim, bias=config.qkv_proj_bias)
748
+
749
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
750
+ self.causal = config.causal
751
+ self.drop = nn.Dropout(config.attn_pdrop)
752
+
753
+ def forward(
754
+ self,
755
+ hidden_states: torch.Tensor,
756
+ attention_mask: Optional[torch.Tensor] = None,
757
+ position_ids: Optional[torch.LongTensor] = None,
758
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
759
+ output_attentions: bool = False,
760
+ use_cache: bool = False,
761
+ is_padded_inputs: Optional[bool] = True,
762
+ cu_seqlens: Optional[torch.Tensor] = None,
763
+ max_seq_len: Optional[int] = None,
764
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
765
+
766
+ has_layer_past = past_key_value is not None
767
+
768
+ if has_layer_past:
769
+ past_key_value = past_key_value[0]
770
+ past_len = past_key_value[1]
771
+ else:
772
+ past_len = 0
773
+
774
+ qkv = self.Wqkv(hidden_states)
775
+ qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
776
+
777
+ past_key_value = (past_key_value, past_len + qkv.size(1)) if use_cache else None
778
+
779
+ if self.rotary_emb_dim > 0:
780
+ if self.rotary_head_dim:
781
+ qkv = rearrange(qkv, "b s three h d -> b h three s d")
782
+ qkv = self.rotary_emb(qkv, seqlen_offset=past_len)
783
+
784
+ if self.rotary_head_dim:
785
+ qkv = rearrange(qkv, "b h three s d -> b s three h d")
786
+
787
+ query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
788
+
789
+ query = query.permute(0, 2, 1, 3)
790
+ key = key.permute(0, 2, 1, 3)
791
+ value = value.permute(0, 2, 1, 3)
792
+
793
+ attention_scores = torch.matmul(query, key.transpose(-1, -2)) / self.norm_factor
794
+ if attention_mask is not None:
795
+ attention_scores = attention_scores + attention_mask
796
+
797
+ attentions_probs = F.softmax(attention_scores, dim=-1)
798
+ attentions_probs = self.drop(attentions_probs)
799
+
800
+ attn_output = torch.matmul(attentions_probs, value)
801
+ attn_output = rearrange(attn_output.permute(0, 2, 1, 3), "... h d -> ... (h d)")
802
+
803
+ attn_output = self.out_proj(attn_output)
804
+
805
+ return attn_output
806
+
807
+
808
+ class NomicBertBlock(nn.Module):
809
+ def __init__(
810
+ self,
811
+ config,
812
+ ):
813
+ super().__init__()
814
+ self.prenorm = config.prenorm
815
+ self.fused_dropout_add_ln = config.fused_dropout_add_ln
816
+
817
+ self.attn = NomicBertAttention(config)
818
+ activation = (
819
+ F.sigmoid
820
+ if config.activation_function == "glu"
821
+ else (F.silu if config.activation_function == "swiglu" else F.gelu)
822
+ )
823
+ if config.activation_function in ["glu", "swiglu", "geglu"]:
824
+ self.mlp = NomciBertGatedMLP(
825
+ config.n_embd,
826
+ hidden_features=config.n_inner,
827
+ bias1=config.mlp_fc1_bias,
828
+ bias2=config.mlp_fc2_bias,
829
+ activation=activation,
830
+ fused_bias_fc=config.fused_bias_fc,
831
+ )
832
+ else:
833
+ self.mlp = NomicBertMLP(
834
+ config.n_embd,
835
+ hidden_features=config.n_inner,
836
+ bias1=config.mlp_fc1_bias,
837
+ bias2=config.mlp_fc2_bias,
838
+ activation=activation,
839
+ fused_bias_fc=config.fused_bias_fc,
840
+ )
841
+
842
+ self.dropout1 = nn.Dropout(config.resid_pdrop)
843
+ self.norm1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
844
+ self.norm2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
845
+ self.dropout2 = nn.Dropout(config.resid_pdrop)
846
+
847
+ def forward(
848
+ self,
849
+ hidden_states: torch.Tensor,
850
+ hidden_states2: torch.Tensor,
851
+ residual: Optional[torch.Tensor] = None,
852
+ attention_mask: Optional[torch.Tensor] = None,
853
+ position_ids: Optional[torch.LongTensor] = None,
854
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
855
+ is_padded_inputs: Optional[bool] = True,
856
+ output_attentions: Optional[bool] = False,
857
+ use_cache: Optional[bool] = False,
858
+ cu_seqlens: Optional[torch.Tensor] = None,
859
+ max_seq_len: Optional[int] = None,
860
+ ):
861
+ r"""Pass the input through the encoder layer.
862
+
863
+ Args:
864
+ hidden_states: the sequence to the encoder layer (required).
865
+ residual: if postnorm, residual=None, If prenorm, hidden_states = Attn/MLP(LN(residual))
866
+ mixer_subset: for cross-attention only. If not None, will take a subset of x
867
+ before applying the query projection. Useful for e.g., ViT where we only care
868
+ about the CLS token in the last layer.
869
+ """
870
+ if self.prenorm:
871
+ dropped = self.dropout1(hidden_states)
872
+ residual = (dropped + residual) if residual is not None else dropped
873
+ hidden_states = self.norm1(residual.to(dtype=self.norm1.weight.dtype))
874
+ hidden_states = self.attn(
875
+ hidden_states,
876
+ attention_mask=attention_mask,
877
+ is_padded_inputs=is_padded_inputs,
878
+ cu_seqlens=cu_seqlens,
879
+ max_seq_len=max_seq_len,
880
+ )
881
+
882
+ dropped = self.dropout2(hidden_states)
883
+ residual = (dropped + residual) if residual is not None else dropped
884
+ hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
885
+ hidden_states = self.mlp(hidden_states)
886
+
887
+ return hidden_states, None, residual
888
+ else:
889
+ assert residual is None
890
+ attn_outputs = self.attn(
891
+ hidden_states,
892
+ attention_mask=attention_mask,
893
+ is_padded_inputs=is_padded_inputs,
894
+ cu_seqlens=cu_seqlens,
895
+ max_seq_len=max_seq_len,
896
+ )
897
+ hidden_states = self.norm1((self.dropout1(attn_outputs) + hidden_states).to(dtype=self.norm1.weight.dtype))
898
+ mlp_out = self.mlp(hidden_states)
899
+
900
+ hidden_states = self.norm2((self.dropout2(mlp_out) + hidden_states).to(dtype=self.norm2.weight.dtype))
901
+ return hidden_states, None, None
902
+
903
+
904
+ class NomicBertEncoder(nn.Module):
905
+ def __init__(self, config: GPT2Config):
906
+ super().__init__()
907
+ self.layers = nn.ModuleList([NomicBertBlock(config) for _ in range(config.n_layer)])
908
+ self.gradient_checkpointing = False
909
+ self.config = config
910
+
911
+ def forward(
912
+ self,
913
+ hidden_states: torch.LongTensor = None,
914
+ attention_mask: Optional[torch.Tensor] = None,
915
+ position_ids: Optional[torch.LongTensor] = None,
916
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
917
+ inputs_embeds: Optional[torch.FloatTensor] = None,
918
+ use_cache: Optional[bool] = None,
919
+ output_attentions: Optional[bool] = None,
920
+ output_hidden_states: Optional[bool] = None,
921
+ return_dict: Optional[bool] = None,
922
+ is_padded_inputs: Optional[bool] = True,
923
+ ):
924
+ """If subset_mask is not None, we only want output for the subset of the sequence.
925
+ This means that we only compute the last layer output for these tokens.
926
+ subset_mask: (batch, seqlen), dtype=torch.bool
927
+ """
928
+ hidden_states2 = None
929
+ residual = None
930
+
931
+ for _, layer in enumerate(self.layers):
932
+ if self.gradient_checkpointing and self.training:
933
+
934
+ def create_custom_forward(module):
935
+ def custom_forward(*inputs):
936
+ # None for past_key_value
937
+ return module(*inputs)
938
+
939
+ return custom_forward
940
+
941
+ hidden_states, hidden_states2, residual = torch.utils.checkpoint.checkpoint(
942
+ create_custom_forward(layer),
943
+ hidden_states,
944
+ hidden_states2,
945
+ residual,
946
+ attention_mask,
947
+ None,
948
+ None,
949
+ is_padded_inputs,
950
+ # if you freeze ANY layers, you need `use_reentrant=False`
951
+ # https://github.com/huggingface/transformers/issues/21381
952
+ # https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/7
953
+ use_reentrant=False,
954
+ )
955
+
956
+ else:
957
+ hidden_states, hidden_states2, residual = layer(
958
+ hidden_states,
959
+ hidden_states2,
960
+ residual,
961
+ attention_mask,
962
+ position_ids,
963
+ None,
964
+ is_padded_inputs,
965
+ output_attentions,
966
+ use_cache,
967
+ )
968
+ return hidden_states
969
+
970
+
971
+ class NomicBertPooler(nn.Module):
972
+ def __init__(self, config):
973
+ super().__init__()
974
+ self.dense = nn.Linear(config.n_embd, config.n_embd)
975
+ self.activation = nn.Tanh()
976
+
977
+ def forward(self, hidden_states, pool=True):
978
+ # We "pool" the model by simply taking the hidden state corresponding
979
+ # to the first token.
980
+ first_token_tensor = hidden_states[:, 0] if pool else hidden_states
981
+ pooled_output = self.dense(first_token_tensor)
982
+ pooled_output = self.activation(pooled_output)
983
+ return pooled_output
984
+
985
+
986
+ class NomicBertPredictionHeadTransform(nn.Module):
987
+ def __init__(self, config):
988
+ super().__init__()
989
+ self.dense = nn.Linear(config.n_embd, config.n_embd, bias=config.mlp_fc1_bias)
990
+ approximate = "tanh" if config.activation_function in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] else "none"
991
+ if config.activation_function == "swiglu":
992
+ self.transform_act_fn = F.silu
993
+ else:
994
+ self.transform_act_fn = nn.GELU(approximate=approximate)
995
+
996
+ self.layer_norm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
997
+
998
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
999
+ hidden_states = self.dense(hidden_states)
1000
+ hidden_states = self.transform_act_fn(hidden_states)
1001
+ hidden_states = self.layer_norm(hidden_states)
1002
+
1003
+ return hidden_states
1004
+
1005
+
1006
+ class NomicBertLMPredictionHead(nn.Module):
1007
+ def __init__(self, config):
1008
+ super().__init__()
1009
+
1010
+ self.transform = NomicBertPredictionHeadTransform(config)
1011
+
1012
+ self.decoder = nn.Linear(config.n_embd, config.vocab_size, bias=config.mlp_fc1_bias)
1013
+
1014
+ def forward(self, hidden_states):
1015
+ hidden_states = self.transform(hidden_states)
1016
+ hidden_states = self.decoder(hidden_states)
1017
+ return hidden_states
1018
+
1019
+
1020
+ class NomicBertPreTrainingHeads(nn.Module):
1021
+ def __init__(self, config):
1022
+ super().__init__()
1023
+ self.predictions = NomicBertLMPredictionHead(config)
1024
+
1025
+ def forward(self, sequence_output):
1026
+ prediction_scores = self.predictions(sequence_output)
1027
+ return prediction_scores
1028
+
1029
+
1030
+ class NomicBertModel(NomicBertPreTrainedModel):
1031
+ def __init__(self, config: GPT2Config, add_pooling_layer=True):
1032
+ super().__init__(config)
1033
+ self.pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
1034
+ if config.vocab_size % self.pad_vocab_size_multiple != 0:
1035
+ config.vocab_size += self.pad_vocab_size_multiple - (config.vocab_size % self.pad_vocab_size_multiple)
1036
+
1037
+ assert config.activation_function in [
1038
+ "gelu",
1039
+ "gelu_new",
1040
+ "gelu_fast",
1041
+ "gelu_pytorch_tanh",
1042
+ "swiglu",
1043
+ "geglu",
1044
+ "glu",
1045
+ ]
1046
+
1047
+ self.embeddings = NomicBertEmbeddings(config)
1048
+ self.emb_drop = nn.Dropout(config.resid_pdrop)
1049
+ self.emb_ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
1050
+ self.encoder = NomicBertEncoder(config)
1051
+ self.pooler = NomicBertPooler(config) if add_pooling_layer else None
1052
+
1053
+ self.apply(partial(_init_weights, initializer_range=config.initializer_range))
1054
+
1055
+ def forward(
1056
+ self,
1057
+ input_ids,
1058
+ position_ids=None,
1059
+ token_type_ids=None,
1060
+ attention_mask=None,
1061
+ return_dict=None,
1062
+ ):
1063
+ if token_type_ids is None:
1064
+ token_type_ids = torch.zeros_like(input_ids)
1065
+ hidden_states = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)
1066
+ hidden_states = self.emb_ln(hidden_states)
1067
+ hidden_states = self.emb_drop(hidden_states)
1068
+
1069
+ attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape)
1070
+ sequence_output = self.encoder(hidden_states, attention_mask=attention_mask, return_dict=return_dict)
1071
+
1072
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
1073
+
1074
+ return BaseModelOutputWithPoolingAndCrossAttentions(
1075
+ last_hidden_state=sequence_output,
1076
+ pooler_output=pooled_output,
1077
+ )
1078
+
1079
+
1080
+ class NomicBertForPreTraining(NomicBertPreTrainedModel):
1081
+ _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
1082
+
1083
+ def __init__(self, config: GPT2Config):
1084
+ super().__init__(config)
1085
+
1086
+ self.bert = NomicBertModel(config, add_pooling_layer=getattr(config, "add_pooling_layer", False))
1087
+ self.cls = NomicBertPreTrainingHeads(config)
1088
+ self.mlm_loss = nn.CrossEntropyLoss()
1089
+
1090
+ # Initialize weights and apply final processing
1091
+ self.apply(partial(_init_weights, initializer_range=config.initializer_range))
1092
+ self.tie_weights()
1093
+
1094
+ def tie_weights(self):
1095
+ self.cls.predictions.decoder.weight = self.bert.embeddings.word_embeddings.weight
1096
+
1097
+ def forward(
1098
+ self,
1099
+ input_ids,
1100
+ position_ids=None,
1101
+ token_type_ids=None,
1102
+ attention_mask=None,
1103
+ labels=None,
1104
+ ):
1105
+ """
1106
+ If labels are provided, they must be -100 for masked out tokens (as specified in the attention
1107
+ mask).
1108
+ Outputs:
1109
+ if `labels` and `next_sentence_label` are not `None`:
1110
+ Outputs the total_loss which is the sum of the masked language modeling loss and the next
1111
+ sentence classification loss.
1112
+ if `labels` or `next_sentence_label` is `None`:
1113
+ Outputs a tuple comprising
1114
+ - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
1115
+ - the next sentence classification logits of shape [batch_size, 2].
1116
+
1117
+ """
1118
+ outputs = self.bert(
1119
+ input_ids,
1120
+ position_ids=position_ids,
1121
+ token_type_ids=token_type_ids,
1122
+ attention_mask=attention_mask.bool() if attention_mask is not None else None,
1123
+ )
1124
+ sequence_output, _ = outputs.last_hidden_state, outputs.pooler_output
1125
+
1126
+ prediction_scores = self.cls(sequence_output)
1127
+
1128
+ total_loss = None
1129
+ if labels is not None:
1130
+ masked_lm_loss = self.mlm_loss(
1131
+ rearrange(prediction_scores, "... v -> (...) v"),
1132
+ rearrange(labels, "... -> (...)"),
1133
+ )
1134
+ total_loss = masked_lm_loss.float()
1135
+
1136
+ return MaskedLMOutput(
1137
+ loss=total_loss,
1138
+ logits=prediction_scores,
1139
+ hidden_states=outputs.hidden_states,
1140
+ attentions=None,
1141
+ )
1142
+
1143
+
1144
+ class NomicBertForSequenceClassification(NomicBertPreTrainedModel):
1145
+ def __init__(self, config):
1146
+ super().__init__(config)
1147
+ self.num_labels = config.num_labels
1148
+ self.config = config
1149
+
1150
+ self.bert = NomicBertModel(config)
1151
+ classifier_dropout = getattr(config, "classifier_dropout", config.embd_pdrop)
1152
+ self.dropout = nn.Dropout(classifier_dropout)
1153
+ self.classifier = nn.Linear(config.n_embd, config.num_labels)
1154
+
1155
+ # Initialize weights and apply final processing
1156
+ self.post_init()
1157
+
1158
+ def forward(
1159
+ self,
1160
+ input_ids: Optional[torch.Tensor] = None,
1161
+ attention_mask: Optional[torch.Tensor] = None,
1162
+ token_type_ids: Optional[torch.Tensor] = None,
1163
+ position_ids: Optional[torch.Tensor] = None,
1164
+ head_mask: Optional[torch.Tensor] = None,
1165
+ inputs_embeds: Optional[torch.Tensor] = None,
1166
+ labels: Optional[torch.Tensor] = None,
1167
+ output_attentions: Optional[bool] = None,
1168
+ output_hidden_states: Optional[bool] = None,
1169
+ return_dict: Optional[bool] = None,
1170
+ ):
1171
+ r"""
1172
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1173
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1174
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1175
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1176
+ """
1177
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1178
+ outputs = self.bert(
1179
+ input_ids,
1180
+ position_ids=position_ids,
1181
+ token_type_ids=token_type_ids,
1182
+ attention_mask=attention_mask.bool() if attention_mask is not None else None,
1183
+ )
1184
+
1185
+ pooled_output = outputs[1]
1186
+
1187
+ pooled_output = self.dropout(pooled_output)
1188
+ logits = self.classifier(pooled_output)
1189
+
1190
+ loss = None
1191
+ if labels is not None:
1192
+ if self.config.problem_type is None:
1193
+ if self.num_labels == 1:
1194
+ self.config.problem_type = "regression"
1195
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1196
+ self.config.problem_type = "single_label_classification"
1197
+ else:
1198
+ self.config.problem_type = "multi_label_classification"
1199
+
1200
+ if self.config.problem_type == "regression":
1201
+ loss_fct = nn.MSELoss()
1202
+ if self.num_labels == 1:
1203
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
1204
+ else:
1205
+ loss = loss_fct(logits, labels)
1206
+ elif self.config.problem_type == "single_label_classification":
1207
+ loss_fct = nn.CrossEntropyLoss()
1208
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1209
+ elif self.config.problem_type == "multi_label_classification":
1210
+ loss_fct = nn.BCEWithLogitsLoss()
1211
+ loss = loss_fct(logits, labels)
1212
+ if not return_dict:
1213
+ output = (logits,) + outputs[2:]
1214
+ return ((loss,) + output) if loss is not None else output
1215
+
1216
+ return SequenceClassifierOutput(
1217
+ loss=loss,
1218
+ logits=logits,
1219
+ hidden_states=outputs.hidden_states,
1220
+ attentions=outputs.attentions,
1221
+ )
modules.json ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "idx": 0,
4
+ "name": "0",
5
+ "path": "",
6
+ "type": "sentence_transformers.models.Transformer"
7
+ },
8
+ {
9
+ "idx": 1,
10
+ "name": "1",
11
+ "path": "1_Pooling",
12
+ "type": "sentence_transformers.models.Pooling"
13
+ },
14
+ {
15
+ "idx": 2,
16
+ "name": "2",
17
+ "path": "2_Normalize",
18
+ "type": "sentence_transformers.models.Normalize"
19
+ }
20
+ ]
sentence_bert_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "max_seq_length": 1024,
3
+ "do_lower_case": false
4
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cls_token": {
3
+ "content": "[CLS]",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "mask_token": {
10
+ "content": "[MASK]",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "[PAD]",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "sep_token": {
24
+ "content": "[SEP]",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ },
30
+ "unk_token": {
31
+ "content": "[UNK]",
32
+ "lstrip": false,
33
+ "normalized": false,
34
+ "rstrip": false,
35
+ "single_word": false
36
+ }
37
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "[PAD]",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ },
11
+ "100": {
12
+ "content": "[UNK]",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false,
17
+ "special": true
18
+ },
19
+ "101": {
20
+ "content": "[CLS]",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false,
25
+ "special": true
26
+ },
27
+ "102": {
28
+ "content": "[SEP]",
29
+ "lstrip": false,
30
+ "normalized": false,
31
+ "rstrip": false,
32
+ "single_word": false,
33
+ "special": true
34
+ },
35
+ "103": {
36
+ "content": "[MASK]",
37
+ "lstrip": false,
38
+ "normalized": false,
39
+ "rstrip": false,
40
+ "single_word": false,
41
+ "special": true
42
+ }
43
+ },
44
+ "clean_up_tokenization_spaces": true,
45
+ "cls_token": "[CLS]",
46
+ "do_lower_case": true,
47
+ "mask_token": "[MASK]",
48
+ "max_length": 1024,
49
+ "model_max_length": 8192,
50
+ "pad_to_multiple_of": null,
51
+ "pad_token": "[PAD]",
52
+ "pad_token_type_id": 0,
53
+ "padding_side": "right",
54
+ "sep_token": "[SEP]",
55
+ "stride": 0,
56
+ "strip_accents": null,
57
+ "tokenize_chinese_chars": true,
58
+ "tokenizer_class": "BertTokenizer",
59
+ "truncation_side": "right",
60
+ "truncation_strategy": "longest_first",
61
+ "unk_token": "[UNK]"
62
+ }
vocab.txt ADDED
The diff for this file is too large to render. See raw diff