full set multi loss ESCI triplets

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+  "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,
+  "include_prompt": true
+}

README.md

@@ -0,0 +1,950 @@
+---
+language: []
+library_name: sentence-transformers
+tags:
+- sentence-transformers
+- sentence-similarity
+- feature-extraction
+- generated_from_trainer
+- dataset_size:1182198
+- loss:CachedMultipleNegativesRankingLoss
+- loss:AnglELoss
+base_model: nomic-ai/nomic-embed-text-v1.5
+datasets: []
+metrics:
+- cosine_accuracy
+- dot_accuracy
+- manhattan_accuracy
+- euclidean_accuracy
+- max_accuracy
+- pearson_cosine
+- spearman_cosine
+- pearson_manhattan
+- spearman_manhattan
+- pearson_euclidean
+- spearman_euclidean
+- pearson_dot
+- spearman_dot
+- pearson_max
+- spearman_max
+widget:
+- source_sentence: dog instrument toy
+  sentences:
+  - VATOS 25-in-1 Mars Rover Building Kit Outer Space Explorer Educational Construction
+    Toy for Kids 556 Pieces Solar Powered STEM Science Building Blocks Set, VATOS,
+    White
+  - Prefer Green 7 PCS Portion Control Containers Kit (with COMPLETE GUIDE & 21 DAY
+    DAILY TRACKER & 21 DAY MEAL PLANNER & RECIPES PDFs),Label-Coded,Multi-Color-Coded
+    System,Perfect Size for Lose Weight, Prefer Green, 7 PCS
+  - Coolibar UPF 50+ Men's Women's Gannett UV Gloves - Sun Protective (Medium- Light
+    Blue), Coolibar, Light Blue
+- source_sentence: flame decal stickers
+  sentences:
+  - Tribal Flames Splash Pair - Vinyl Decal Sticker - 12" x 5" - Blue Flames, Sticker
+    Pimp, Blue Flames
+  - PC Gaming Headset Headphone Hook Holder Hanger Mount, Headphones Stand with Adjustable
+    & Rotating Arm Clamp , Under Desk Design , Universal Fit , Built in Cable Clip
+    Organizer EURPMASK, EURPMASK Choose the color of europe, Black
+  - Quick Charge 3.0 Wall Charger, 4-Pack 18W QC 3.0 USB Charger Adapter Fast Charging
+    Block Compatible Wireless Charger Compatible with Samsung Galaxy S10 S9 S8 Plus
+    S7 S6 Edge Note 9, LG, Kindle, Tablet, HONOT, Black
+- source_sentence: 'search_query: softies women''s ultra soft marshmallow hooded lounger'
+  sentences:
+  - 'search_document: Red-A Placemats for Dining Table Set of 6 Heat-Resistant Wipeable
+    Table Mats for Kitchen Table Decoration Waterproof Vinyl Placemats Easy to Clean,Black
+    w/Brown, Red-A, Black'
+  - 'search_document: Softies Women''s Ultra Soft Marshmallow Hooded Lounger, Platinum,
+    L/XL, Softies, Platinum'
+  - 'search_document: Ekouaer Women''s Sleepwear Robe with Pockets Plus Size Maxi
+    Lounger Zipper Short Sleeve Bathrobe Housecoat (Black,L), Ekouaer, Black'
+- source_sentence: 'search_query: wine glasses without stem'
+  sentences:
+  - 'search_document: STAUBER Best Bulb Changer with PowerLatch Extension Pole (Large
+    Suction, 4 Feet), STAUBER, Large Suction'
+  - 'search_document: Hand Blown Italian Style Crystal Burgundy Wine Glasses - Lead-Free
+    Premium Crystal Clear Glass - Set of 2 - 21 Ounce - Gift-Box for any Occasion,
+    JBHO, Burgundy'
+  - 'search_document: MyGift Modern Copper Stemless Wine Glasses, Set of 4, MyGift,
+    Copper'
+- source_sentence: 'search_query: weighted blanket without glass beads'
+  sentences:
+  - 'search_document: Eigso Women Men Spike Punk Rock Black Leather Cuff Rivet Bracelet
+    Bangle Adjustable Snap Button, Eigso, Black'
+  - 'search_document: Quility Weighted Blanket with Soft Cover - 20 lbs Full/Queen
+    Size Heavy Blanket for Adults - Heating & Cooling, Machine Washable - (60" X 80")
+    (Navy), Quility, Navy Cover + Grey Cotton Blanket'
+  - 'search_document: Bedsure Queen Weighted Blanket 15 Pounds - Adult Weighted Blanket
+    60x80 - Soft Heavy Blanket with Breathable TPE Insert No Glass Beads, Bedsure,
+    Navy'
+pipeline_tag: sentence-similarity
+model-index:
+- name: SentenceTransformer based on nomic-ai/nomic-embed-text-v1.5
+  results:
+  - task:
+      type: triplet
+      name: Triplet
+    dataset:
+      name: Unknown
+      type: unknown
+    metrics:
+    - type: cosine_accuracy
+      value: 0.7236
+      name: Cosine Accuracy
+    - type: dot_accuracy
+      value: 0.282
+      name: Dot Accuracy
+    - type: manhattan_accuracy
+      value: 0.7231
+      name: Manhattan Accuracy
+    - type: euclidean_accuracy
+      value: 0.7227
+      name: Euclidean Accuracy
+    - type: max_accuracy
+      value: 0.7236
+      name: Max Accuracy
+  - task:
+      type: semantic-similarity
+      name: Semantic Similarity
+    dataset:
+      name: Unknown
+      type: unknown
+    metrics:
+    - type: pearson_cosine
+      value: 0.4912162846043421
+      name: Pearson Cosine
+    - type: spearman_cosine
+      value: 0.4658522123059972
+      name: Spearman Cosine
+    - type: pearson_manhattan
+      value: 0.4599741171303018
+      name: Pearson Manhattan
+    - type: spearman_manhattan
+      value: 0.4428141949345816
+      name: Spearman Manhattan
+    - type: pearson_euclidean
+      value: 0.46194545823984606
+      name: Pearson Euclidean
+    - type: spearman_euclidean
+      value: 0.44478471500226807
+      name: Spearman Euclidean
+    - type: pearson_dot
+      value: 0.45451995456560107
+      name: Pearson Dot
+    - type: spearman_dot
+      value: 0.43844636325741904
+      name: Spearman Dot
+    - type: pearson_max
+      value: 0.4912162846043421
+      name: Pearson Max
+    - type: spearman_max
+      value: 0.4658522123059972
+      name: Spearman Max
+---
+
+# SentenceTransformer based on nomic-ai/nomic-embed-text-v1.5
+
+This is a [sentence-transformers](https://www.SBERT.net) model finetuned from [nomic-ai/nomic-embed-text-v1.5](https://huggingface.co/nomic-ai/nomic-embed-text-v1.5) on the triplets and pairs datasets. It maps sentences & paragraphs to a 768-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more.
+
+## Model Details
+
+### Model Description
+- **Model Type:** Sentence Transformer
+- **Base model:** [nomic-ai/nomic-embed-text-v1.5](https://huggingface.co/nomic-ai/nomic-embed-text-v1.5) <!-- at revision b0753ae76394dd36bcfb912a46018088bca48be0 -->
+- **Maximum Sequence Length:** 8192 tokens
+- **Output Dimensionality:** 768 tokens
+- **Similarity Function:** Cosine Similarity
+- **Training Datasets:**
+    - triplets
+    - pairs
+<!-- - **Language:** Unknown -->
+<!-- - **License:** Unknown -->
+
+### Model Sources
+
+- **Documentation:** [Sentence Transformers Documentation](https://sbert.net)
+- **Repository:** [Sentence Transformers on GitHub](https://github.com/UKPLab/sentence-transformers)
+- **Hugging Face:** [Sentence Transformers on Hugging Face](https://huggingface.co/models?library=sentence-transformers)
+
+### Full Model Architecture
+
+```
+SentenceTransformer(
+  (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False}) with Transformer model: NomicBertModel 
+  (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, 'include_prompt': True})
+)
+```
+
+## Usage
+
+### Direct Usage (Sentence Transformers)
+
+First install the Sentence Transformers library:
+
+```bash
+pip install -U sentence-transformers
+```
+
+Then you can load this model and run inference.
+```python
+from sentence_transformers import SentenceTransformer
+
+# Download from the 🤗 Hub
+model = SentenceTransformer("lv12/esci-nomic-embed-text-v1_5_4")
+# Run inference
+sentences = [
+    'search_query: weighted blanket without glass beads',
+    'search_document: Bedsure Queen Weighted Blanket 15 Pounds - Adult Weighted Blanket 60x80 - Soft Heavy Blanket with Breathable TPE Insert No Glass Beads, Bedsure, Navy',
+    'search_document: Quility Weighted Blanket with Soft Cover - 20 lbs Full/Queen Size Heavy Blanket for Adults - Heating & Cooling, Machine Washable - (60" X 80") (Navy), Quility, Navy Cover + Grey Cotton Blanket',
+]
+embeddings = model.encode(sentences)
+print(embeddings.shape)
+# [3, 768]
+
+# Get the similarity scores for the embeddings
+similarities = model.similarity(embeddings, embeddings)
+print(similarities.shape)
+# [3, 3]
+```
+
+<!--
+### Direct Usage (Transformers)
+
+<details><summary>Click to see the direct usage in Transformers</summary>
+
+</details>
+-->
+
+<!--
+### Downstream Usage (Sentence Transformers)
+
+You can finetune this model on your own dataset.
+
+<details><summary>Click to expand</summary>
+
+</details>
+-->
+
+<!--
+### Out-of-Scope Use
+
+*List how the model may foreseeably be misused and address what users ought not to do with the model.*
+-->
+
+## Evaluation
+
+### Metrics
+
+#### Triplet
+
+* Evaluated with [<code>TripletEvaluator</code>](https://sbert.net/docs/package_reference/sentence_transformer/evaluation.html#sentence_transformers.evaluation.TripletEvaluator)
+
+| Metric              | Value      |
+|:--------------------|:-----------|
+| **cosine_accuracy** | **0.7236** |
+| dot_accuracy        | 0.282      |
+| manhattan_accuracy  | 0.7231     |
+| euclidean_accuracy  | 0.7227     |
+| max_accuracy        | 0.7236     |
+
+#### Semantic Similarity
+
+* Evaluated with [<code>EmbeddingSimilarityEvaluator</code>](https://sbert.net/docs/package_reference/sentence_transformer/evaluation.html#sentence_transformers.evaluation.EmbeddingSimilarityEvaluator)
+
+| Metric              | Value      |
+|:--------------------|:-----------|
+| pearson_cosine      | 0.4912     |
+| **spearman_cosine** | **0.4659** |
+| pearson_manhattan   | 0.46       |
+| spearman_manhattan  | 0.4428     |
+| pearson_euclidean   | 0.4619     |
+| spearman_euclidean  | 0.4448     |
+| pearson_dot         | 0.4545     |
+| spearman_dot        | 0.4384     |
+| pearson_max         | 0.4912     |
+| spearman_max        | 0.4659     |
+
+<!--
+## Bias, Risks and Limitations
+
+*What are the known or foreseeable issues stemming from this model? You could also flag here known failure cases or weaknesses of the model.*
+-->
+
+<!--
+### Recommendations
+
+*What are recommendations with respect to the foreseeable issues? For example, filtering explicit content.*
+-->
+
+## Training Details
+
+### Training Datasets
+
+#### triplets
+
+* Dataset: triplets
+* Size: 684,084 training samples
+* Columns: <code>anchor</code>, <code>positive</code>, and <code>negative</code>
+* Approximate statistics based on the first 1000 samples:
+  |         | anchor                                                                           | positive                                                                           | negative                                                                           |
+  |:--------|:---------------------------------------------------------------------------------|:-----------------------------------------------------------------------------------|:-----------------------------------------------------------------------------------|
+  | type    | string                                                                           | string                                                                             | string                                                                             |
+  | details | <ul><li>min: 7 tokens</li><li>mean: 11.1 tokens</li><li>max: 22 tokens</li></ul> | <ul><li>min: 17 tokens</li><li>mean: 42.75 tokens</li><li>max: 95 tokens</li></ul> | <ul><li>min: 15 tokens</li><li>mean: 43.8 tokens</li><li>max: 127 tokens</li></ul> |
+* Samples:
+  | anchor                                                          | positive                                                                                                                                                                                                                                          | negative                                                                                                                                                                                                                                        |
+  |:----------------------------------------------------------------|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+  | <code>search_query: tarps heavy duty waterproof 8x10</code>     | <code>search_document: 8' x 10' Super Heavy Duty 16 Mil Brown Poly Tarp Cover - Thick Waterproof, UV Resistant, Rip and Tear Proof Tarpaulin with Grommets and Reinforced Edges - by Xpose Safety, Xpose Safety, Brown</code>                     | <code>search_document: Grillkid 6'X8' 4.5 Mil Thick General Purpose Waterproof Poly Tarp, Grillkid, All Purpose</code>                                                                                                                          |
+  | <code>search_query: wireless keyboard without number pad</code> | <code>search_document: Macally 2.4G Small Wireless Keyboard - Ergonomic & Comfortable Computer Keyboard - Compact Keyboard for Laptop or Windows PC Desktop, Tablet, Smart TV - Plug & Play Mini Keyboard with 12 Hot Keys, Macally, Black</code> | <code>search_document: Wireless Keyboard - iClever GKA22S Rechargeable Keyboard with Number Pad, Full-Size Stainless Steel Ultra Slim Keyboard, 2.4G Stable Connection Wireless Keyboard for iMac, Mackbook, PC, Laptop, iClever, Silver</code> |
+  | <code>search_query: geometry earrings</code>                    | <code>search_document: Simple Stud Earrings for Women, Geometric Minimalist Stud Earring Set Tiny Circle Triangle Square Bar Stud Earrings Mini Cartilage Tragus Earrings, choice of all, B:Circle Sliver</code>                                  | <code>search_document: BONALUNA Bohemian Wood And Marble Effect Oblong Shaped Drop Statement Earrings (VIVID TURQUOISE), BONALUNA, VIVID TURQUOISE</code>                                                                                       |
+* Loss: [<code>CachedMultipleNegativesRankingLoss</code>](https://sbert.net/docs/package_reference/sentence_transformer/losses.html#cachedmultiplenegativesrankingloss) with these parameters:
+  ```json
+  {
+      "scale": 20.0,
+      "similarity_fct": "cos_sim"
+  }
+  ```
+
+#### pairs
+
+* Dataset: pairs
+* Size: 498,114 training samples
+* Columns: <code>sentence1</code>, <code>sentence2</code>, and <code>score</code>
+* Approximate statistics based on the first 1000 samples:
+  |         | sentence1                                                                        | sentence2                                                                          | score                                                          |
+  |:--------|:---------------------------------------------------------------------------------|:-----------------------------------------------------------------------------------|:---------------------------------------------------------------|
+  | type    | string                                                                           | string                                                                             | float                                                          |
+  | details | <ul><li>min: 3 tokens</li><li>mean: 6.73 tokens</li><li>max: 33 tokens</li></ul> | <ul><li>min: 10 tokens</li><li>mean: 40.14 tokens</li><li>max: 98 tokens</li></ul> | <ul><li>min: 0.0</li><li>mean: 0.81</li><li>max: 1.0</li></ul> |
+* Samples:
+  | sentence1                                                                                                                                                    | sentence2                                                                                                                                                                                                   | score            |
+  |:-------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:-----------------|
+  | <code>I would choose a medium weight waterproof  fabric, hip length jacket or longer, long sleeves, zip front, with a hood and deep pockets with zips</code> | <code>ZSHOW Men's Winter Hooded Packable Down Jacket(Blue, XX-Large), ZSHOW, Blue</code>                                                                                                                    | <code>1.0</code> |
+  | <code>sequin dance costume girls</code>                                                                                                                      | <code>Yeahdor Big Girls' Lyrical Latin Ballet Dance Costumes Dresses Halter Sequins Irregular Tutu Skirted Leotard Dancewear Pink 12-14, Yeahdor, Pink</code>                                               | <code>1.0</code> |
+  | <code>paint easel bulk</code>                                                                                                                                | <code>Artecho Artist Easel Display Easel Stand, 2 Pack Metal Tripod Stand Easel for Painting, Hold Canvas from 21" to 66", Floor and Tabletop Displaying, Painting with Portable Bag, Artecho, Black</code> | <code>1.0</code> |
+* Loss: [<code>AnglELoss</code>](https://sbert.net/docs/package_reference/sentence_transformer/losses.html#angleloss) with these parameters:
+  ```json
+  {
+      "scale": 20.0,
+      "similarity_fct": "pairwise_angle_sim"
+  }
+  ```
+
+### Evaluation Datasets
+
+#### triplets
+
+* Dataset: triplets
+* Size: 10,000 evaluation samples
+* Columns: <code>anchor</code>, <code>positive</code>, and <code>negative</code>
+* Approximate statistics based on the first 1000 samples:
+  |         | anchor                                                                            | positive                                                                            | negative                                                                           |
+  |:--------|:----------------------------------------------------------------------------------|:------------------------------------------------------------------------------------|:-----------------------------------------------------------------------------------|
+  | type    | string                                                                            | string                                                                              | string                                                                             |
+  | details | <ul><li>min: 7 tokens</li><li>mean: 11.13 tokens</li><li>max: 23 tokens</li></ul> | <ul><li>min: 15 tokens</li><li>mean: 43.11 tokens</li><li>max: 107 tokens</li></ul> | <ul><li>min: 15 tokens</li><li>mean: 43.56 tokens</li><li>max: 99 tokens</li></ul> |
+* Samples:
+  | anchor                                                             | positive                                                                                                                                                                                                                         | negative                                                                                                                                                                                                              |
+  |:-------------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+  | <code>search_query: hitch fifth wheel</code>                       | <code>search_document: ENIXWILL 5th Wheel Trailer Hitch Lifting Device Bracket Pin Fit for Hitch Companion and Patriot Series Hitch, ENIXWILL, Black</code>                                                                      | <code>search_document: ECOTRIC Fifth 5th Wheel Trailer Hitch Mount Rails and Installation Kits for Full-Size Trucks, ECOTRIC, black</code>                                                                            |
+  | <code>search_query: dek pro</code>                                 | <code>search_document: Cubiker Computer Desk 47 inch Home Office Writing Study Desk, Modern Simple Style Laptop Table with Storage Bag, Brown, Cubiker, Brown</code>                                                             | <code>search_document: FEZIBO Dual Motor L Shaped Electric Standing Desk, 48 Inches Stand Up Corner Desk, Home Office Sit Stand Desk with Rustic Brown Top and Black Frame, FEZIBO, Rustic Brown</code>               |
+  | <code>search_query: 1 year baby mouth without teeth cleaner</code> | <code>search_document: Baby Toothbrush,Infant Toothbrush,Baby Tongue Cleaner,Infant Toothbrush,Baby Tongue Cleaner Newborn,Toothbrush Tongue Cleaner Dental Care for 0-36 Month Baby,36 Pcs + Free 4 Pcs, Babycolor, Blue</code> | <code>search_document: Slotic Baby Toothbrush for 0-2 Years, Safe and Sturdy, Toddler Oral Care Teether Brush, Extra Soft Bristle for Baby Teeth and Infant Gums, Dentist Recommended (4-Pack), Slotic, 4 Pack</code> |
+* Loss: [<code>CachedMultipleNegativesRankingLoss</code>](https://sbert.net/docs/package_reference/sentence_transformer/losses.html#cachedmultiplenegativesrankingloss) with these parameters:
+  ```json
+  {
+      "scale": 20.0,
+      "similarity_fct": "cos_sim"
+  }
+  ```
+
+#### pairs
+
+* Dataset: pairs
+* Size: 10,000 evaluation samples
+* Columns: <code>sentence1</code>, <code>sentence2</code>, and <code>score</code>
+* Approximate statistics based on the first 1000 samples:
+  |         | sentence1                                                                       | sentence2                                                                         | score                                                          |
+  |:--------|:--------------------------------------------------------------------------------|:----------------------------------------------------------------------------------|:---------------------------------------------------------------|
+  | type    | string                                                                          | string                                                                            | float                                                          |
+  | details | <ul><li>min: 3 tokens</li><li>mean: 6.8 tokens</li><li>max: 34 tokens</li></ul> | <ul><li>min: 9 tokens</li><li>mean: 39.7 tokens</li><li>max: 101 tokens</li></ul> | <ul><li>min: 0.0</li><li>mean: 0.77</li><li>max: 1.0</li></ul> |
+* Samples:
+  | sentence1                                       | sentence2                                                                                                                                                                                     | score            |
+  |:------------------------------------------------|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:-----------------|
+  | <code>outdoor ceiling fans without light</code> | <code>44" Plaza Industrial Indoor Outdoor Ceiling Fan with Remote Control Oil Rubbed Bronze Damp Rated for Patio Porch - Casa Vieja, Casa Vieja, No Light Kit - Bronze</code>                 | <code>1.0</code> |
+  | <code>bathroom cabinet</code>                   | <code>Homfa Bathroom Floor Cabinet Free Standing with Single Door Multifunctional Bathroom Storage Organizer Toiletries(Ivory White), Homfa, White</code>                                     | <code>1.0</code> |
+  | <code>fitbit charge 3</code>                    | <code>TreasureMax Compatible with Fitbit Charge 2 Bands for Women/Men,Silicone Fadeless Pattern Printed Replacement Floral Bands for Fitbit Charge 2 HR Wristbands, TreasureMax, Paw 2</code> | <code>0.4</code> |
+* Loss: [<code>AnglELoss</code>](https://sbert.net/docs/package_reference/sentence_transformer/losses.html#angleloss) with these parameters:
+  ```json
+  {
+      "scale": 20.0,
+      "similarity_fct": "pairwise_angle_sim"
+  }
+  ```
+
+### Training Hyperparameters
+#### Non-Default Hyperparameters
+
+- `per_device_train_batch_size`: 16
+- `per_device_eval_batch_size`: 4
+- `gradient_accumulation_steps`: 2
+- `learning_rate`: 1e-06
+- `lr_scheduler_type`: cosine_with_restarts
+- `lr_scheduler_kwargs`: {'num_cycles': 1}
+- `warmup_ratio`: 0.01
+- `dataloader_drop_last`: True
+- `dataloader_num_workers`: 4
+- `dataloader_prefetch_factor`: 4
+- `load_best_model_at_end`: True
+- `gradient_checkpointing`: True
+- `batch_sampler`: no_duplicates
+
+#### All Hyperparameters
+<details><summary>Click to expand</summary>
+
+- `overwrite_output_dir`: False
+- `do_predict`: False
+- `prediction_loss_only`: True
+- `per_device_train_batch_size`: 16
+- `per_device_eval_batch_size`: 4
+- `per_gpu_train_batch_size`: None
+- `per_gpu_eval_batch_size`: None
+- `gradient_accumulation_steps`: 2
+- `eval_accumulation_steps`: None
+- `learning_rate`: 1e-06
+- `weight_decay`: 0.0
+- `adam_beta1`: 0.9
+- `adam_beta2`: 0.999
+- `adam_epsilon`: 1e-08
+- `max_grad_norm`: 1.0
+- `num_train_epochs`: 3
+- `max_steps`: -1
+- `lr_scheduler_type`: cosine_with_restarts
+- `lr_scheduler_kwargs`: {'num_cycles': 1}
+- `warmup_ratio`: 0.01
+- `warmup_steps`: 0
+- `log_level`: passive
+- `log_level_replica`: warning
+- `log_on_each_node`: True
+- `logging_nan_inf_filter`: True
+- `save_safetensors`: True
+- `save_on_each_node`: False
+- `save_only_model`: False
+- `no_cuda`: False
+- `use_cpu`: False
+- `use_mps_device`: False
+- `seed`: 42
+- `data_seed`: None
+- `jit_mode_eval`: False
+- `use_ipex`: False
+- `bf16`: False
+- `fp16`: False
+- `fp16_opt_level`: O1
+- `half_precision_backend`: auto
+- `bf16_full_eval`: False
+- `fp16_full_eval`: False
+- `tf32`: None
+- `local_rank`: 0
+- `ddp_backend`: None
+- `tpu_num_cores`: None
+- `tpu_metrics_debug`: False
+- `debug`: []
+- `dataloader_drop_last`: True
+- `dataloader_num_workers`: 4
+- `dataloader_prefetch_factor`: 4
+- `past_index`: -1
+- `disable_tqdm`: False
+- `remove_unused_columns`: True
+- `label_names`: None
+- `load_best_model_at_end`: True
+- `ignore_data_skip`: False
+- `fsdp`: []
+- `fsdp_min_num_params`: 0
+- `fsdp_config`: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}
+- `fsdp_transformer_layer_cls_to_wrap`: None
+- `accelerator_config`: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True}
+- `deepspeed`: None
+- `label_smoothing_factor`: 0.0
+- `optim`: adamw_torch
+- `optim_args`: None
+- `adafactor`: False
+- `group_by_length`: False
+- `length_column_name`: length
+- `ddp_find_unused_parameters`: None
+- `ddp_bucket_cap_mb`: None
+- `ddp_broadcast_buffers`: False
+- `dataloader_pin_memory`: True
+- `dataloader_persistent_workers`: False
+- `skip_memory_metrics`: True
+- `use_legacy_prediction_loop`: False
+- `push_to_hub`: False
+- `resume_from_checkpoint`: None
+- `hub_model_id`: None
+- `hub_strategy`: every_save
+- `hub_private_repo`: False
+- `hub_always_push`: False
+- `gradient_checkpointing`: True
+- `gradient_checkpointing_kwargs`: None
+- `include_inputs_for_metrics`: False
+- `fp16_backend`: auto
+- `push_to_hub_model_id`: None
+- `push_to_hub_organization`: None
+- `mp_parameters`: 
+- `auto_find_batch_size`: False
+- `full_determinism`: False
+- `torchdynamo`: None
+- `ray_scope`: last
+- `ddp_timeout`: 1800
+- `torch_compile`: False
+- `torch_compile_backend`: None
+- `torch_compile_mode`: None
+- `dispatch_batches`: None
+- `split_batches`: None
+- `include_tokens_per_second`: False
+- `include_num_input_tokens_seen`: False
+- `neftune_noise_alpha`: None
+- `batch_sampler`: no_duplicates
+- `multi_dataset_batch_sampler`: proportional
+
+</details>
+
+### Training Logs
+<details><summary>Click to expand</summary>
+
+| Epoch  | Step  | Training Loss | pairs loss | triplets loss | cosine_accuracy | spearman_cosine |
+|:------:|:-----:|:-------------:|:----------:|:-------------:|:---------------:|:---------------:|
+| 0.0027 | 100   | 2.4909        | -          | -             | -               | -               |
+| 0.0054 | 200   | 2.6666        | -          | -             | -               | -               |
+| 0.0081 | 300   | 2.76          | -          | -             | -               | -               |
+| 0.0108 | 400   | 2.6945        | -          | -             | -               | -               |
+| 0.0135 | 500   | 2.9113        | -          | -             | -               | -               |
+| 0.0162 | 600   | 2.3476        | -          | -             | -               | -               |
+| 0.0189 | 700   | 2.2818        | -          | -             | -               | -               |
+| 0.0217 | 800   | 2.4241        | -          | -             | -               | -               |
+| 0.0244 | 900   | 2.5126        | -          | -             | -               | -               |
+| 0.0271 | 1000  | 2.4106        | 4.7376     | 0.8087        | 0.6993          | 0.3844          |
+| 0.0298 | 1100  | 2.2369        | -          | -             | -               | -               |
+| 0.0325 | 1200  | 2.0614        | -          | -             | -               | -               |
+| 0.0352 | 1300  | 2.2178        | -          | -             | -               | -               |
+| 0.0379 | 1400  | 1.974         | -          | -             | -               | -               |
+| 0.0406 | 1500  | 1.9364        | -          | -             | -               | -               |
+| 0.0433 | 1600  | 2.0906        | -          | -             | -               | -               |
+| 0.0460 | 1700  | 1.8783        | -          | -             | -               | -               |
+| 0.0487 | 1800  | 2.1149        | -          | -             | -               | -               |
+| 0.0514 | 1900  | 1.7162        | -          | -             | -               | -               |
+| 0.0541 | 2000  | 1.6761        | 3.8862     | 0.7490        | 0.7097          | 0.4082          |
+| 0.0568 | 2100  | 2.1701        | -          | -             | -               | -               |
+| 0.0596 | 2200  | 2.1306        | -          | -             | -               | -               |
+| 0.0623 | 2300  | 1.6543        | -          | -             | -               | -               |
+| 0.0650 | 2400  | 1.8157        | -          | -             | -               | -               |
+| 0.0677 | 2500  | 1.7779        | -          | -             | -               | -               |
+| 0.0704 | 2600  | 1.9434        | -          | -             | -               | -               |
+| 0.0731 | 2700  | 1.7776        | -          | -             | -               | -               |
+| 0.0758 | 2800  | 1.8197        | -          | -             | -               | -               |
+| 0.0785 | 2900  | 1.9886        | -          | -             | -               | -               |
+| 0.0812 | 3000  | 2.0699        | 3.8031     | 0.7298        | 0.7147          | 0.4282          |
+| 0.0839 | 3100  | 1.9496        | -          | -             | -               | -               |
+| 0.0866 | 3200  | 1.8349        | -          | -             | -               | -               |
+| 0.0893 | 3300  | 2.111         | -          | -             | -               | -               |
+| 0.0920 | 3400  | 1.9956        | -          | -             | -               | -               |
+| 0.0947 | 3500  | 2.0379        | -          | -             | -               | -               |
+| 0.0974 | 3600  | 1.8975        | -          | -             | -               | -               |
+| 0.1002 | 3700  | 1.8552        | -          | -             | -               | -               |
+| 0.1029 | 3800  | 1.9566        | -          | -             | -               | -               |
+| 0.1056 | 3900  | 2.011         | -          | -             | -               | -               |
+| 0.1083 | 4000  | 2.1263        | 3.7799     | 0.7221        | 0.7176          | 0.4393          |
+| 0.1110 | 4100  | 1.8217        | -          | -             | -               | -               |
+| 0.1137 | 4200  | 1.8638        | -          | -             | -               | -               |
+| 0.1164 | 4300  | 1.7699        | -          | -             | -               | -               |
+| 0.1191 | 4400  | 1.8248        | -          | -             | -               | -               |
+| 0.1218 | 4500  | 1.835         | -          | -             | -               | -               |
+| 0.1245 | 4600  | 1.9294        | -          | -             | -               | -               |
+| 0.1272 | 4700  | 1.9817        | -          | -             | -               | -               |
+| 0.1299 | 4800  | 1.877         | -          | -             | -               | -               |
+| 0.1326 | 4900  | 1.5824        | -          | -             | -               | -               |
+| 0.1353 | 5000  | 1.7429        | 3.7728     | 0.7163        | 0.7196          | 0.4496          |
+| 0.1380 | 5100  | 1.8552        | -          | -             | -               | -               |
+| 0.1408 | 5200  | 1.6888        | -          | -             | -               | -               |
+| 0.1435 | 5300  | 1.9409        | -          | -             | -               | -               |
+| 0.1462 | 5400  | 1.9389        | -          | -             | -               | -               |
+| 0.1489 | 5500  | 1.82          | -          | -             | -               | -               |
+| 0.1516 | 5600  | 1.9763        | -          | -             | -               | -               |
+| 0.1543 | 5700  | 1.8122        | -          | -             | -               | -               |
+| 0.1570 | 5800  | 1.7204        | -          | -             | -               | -               |
+| 0.1597 | 5900  | 1.6901        | -          | -             | -               | -               |
+| 0.1624 | 6000  | 1.7785        | 3.7514     | 0.7124        | 0.7195          | 0.4516          |
+| 0.1651 | 6100  | 1.8559        | -          | -             | -               | -               |
+| 0.1678 | 6200  | 1.7646        | -          | -             | -               | -               |
+| 0.1705 | 6300  | 1.9068        | -          | -             | -               | -               |
+| 0.1732 | 6400  | 1.8848        | -          | -             | -               | -               |
+| 0.1759 | 6500  | 1.9384        | -          | -             | -               | -               |
+| 0.1787 | 6600  | 1.7692        | -          | -             | -               | -               |
+| 0.1814 | 6700  | 1.7093        | -          | -             | -               | -               |
+| 0.1841 | 6800  | 1.8759        | -          | -             | -               | -               |
+| 0.1868 | 6900  | 1.7319        | -          | -             | -               | -               |
+| 0.1895 | 7000  | 1.9428        | 3.7487     | 0.7076        | 0.7256          | 0.4496          |
+| 0.1922 | 7100  | 1.5733        | -          | -             | -               | -               |
+| 0.1949 | 7200  | 1.8487        | -          | -             | -               | -               |
+| 0.1976 | 7300  | 1.8361        | -          | -             | -               | -               |
+| 0.2003 | 7400  | 1.9911        | -          | -             | -               | -               |
+| 0.2030 | 7500  | 1.784         | -          | -             | -               | -               |
+| 0.2057 | 7600  | 1.8518        | -          | -             | -               | -               |
+| 0.2084 | 7700  | 1.6232        | -          | -             | -               | -               |
+| 0.2111 | 7800  | 1.6239        | -          | -             | -               | -               |
+| 0.2138 | 7900  | 1.7589        | -          | -             | -               | -               |
+| 0.2165 | 8000  | 1.8644        | 3.7387     | 0.7040        | 0.7241          | 0.4552          |
+| 0.2193 | 8100  | 1.7903        | -          | -             | -               | -               |
+| 0.2220 | 8200  | 1.7197        | -          | -             | -               | -               |
+| 0.2247 | 8300  | 1.9099        | -          | -             | -               | -               |
+| 0.2274 | 8400  | 1.6778        | -          | -             | -               | -               |
+| 0.2301 | 8500  | 1.9249        | -          | -             | -               | -               |
+| 0.2328 | 8600  | 1.8483        | -          | -             | -               | -               |
+| 0.2355 | 8700  | 1.6849        | -          | -             | -               | -               |
+| 0.2382 | 8800  | 1.8647        | -          | -             | -               | -               |
+| 0.2409 | 8900  | 1.8826        | -          | -             | -               | -               |
+| 0.2436 | 9000  | 1.7632        | 3.7403     | 0.7033        | 0.7225          | 0.4545          |
+| 0.2463 | 9100  | 1.8142        | -          | -             | -               | -               |
+| 0.2490 | 9200  | 1.7374        | -          | -             | -               | -               |
+| 0.2517 | 9300  | 1.8646        | -          | -             | -               | -               |
+| 0.2544 | 9400  | 1.7623        | -          | -             | -               | -               |
+| 0.2571 | 9500  | 1.7802        | -          | -             | -               | -               |
+| 0.2599 | 9600  | 1.843         | -          | -             | -               | -               |
+| 0.2626 | 9700  | 1.9797        | -          | -             | -               | -               |
+| 0.2653 | 9800  | 1.7748        | -          | -             | -               | -               |
+| 0.2680 | 9900  | 1.7031        | -          | -             | -               | -               |
+| 0.2707 | 10000 | 1.5536        | 3.7613     | 0.7016        | 0.7259          | 0.4548          |
+| 0.2734 | 10100 | 1.7663        | -          | -             | -               | -               |
+| 0.2761 | 10200 | 1.8218        | -          | -             | -               | -               |
+| 0.2788 | 10300 | 1.6327        | -          | -             | -               | -               |
+| 0.2815 | 10400 | 1.8802        | -          | -             | -               | -               |
+| 0.2842 | 10500 | 1.6294        | -          | -             | -               | -               |
+| 0.2869 | 10600 | 1.9001        | -          | -             | -               | -               |
+| 0.2896 | 10700 | 1.7873        | -          | -             | -               | -               |
+| 0.2923 | 10800 | 1.8121        | -          | -             | -               | -               |
+| 0.2950 | 10900 | 2.0197        | -          | -             | -               | -               |
+| 0.2978 | 11000 | 1.7006        | 3.7559     | 0.7004        | 0.727           | 0.4613          |
+| 0.3005 | 11100 | 1.6404        | -          | -             | -               | -               |
+| 0.3032 | 11200 | 1.9422        | -          | -             | -               | -               |
+| 0.3059 | 11300 | 1.5917        | -          | -             | -               | -               |
+| 0.3086 | 11400 | 1.7236        | -          | -             | -               | -               |
+| 0.3113 | 11500 | 1.8977        | -          | -             | -               | -               |
+| 0.3140 | 11600 | 1.7686        | -          | -             | -               | -               |
+| 0.3167 | 11700 | 1.4493        | -          | -             | -               | -               |
+| 0.3194 | 11800 | 1.7447        | -          | -             | -               | -               |
+| 0.3221 | 11900 | 1.9412        | -          | -             | -               | -               |
+| 0.3248 | 12000 | 1.8           | 3.7308     | 0.6997        | 0.7241          | 0.4618          |
+| 0.3275 | 12100 | 1.8855        | -          | -             | -               | -               |
+| 0.3302 | 12200 | 1.5133        | -          | -             | -               | -               |
+| 0.3329 | 12300 | 1.7893        | -          | -             | -               | -               |
+| 0.3356 | 12400 | 1.7861        | -          | -             | -               | -               |
+| 0.3384 | 12500 | 1.7733        | -          | -             | -               | -               |
+| 0.3411 | 12600 | 1.5877        | -          | -             | -               | -               |
+| 0.3438 | 12700 | 2.03          | -          | -             | -               | -               |
+| 0.3465 | 12800 | 1.7071        | -          | -             | -               | -               |
+| 0.3492 | 12900 | 1.7848        | -          | -             | -               | -               |
+| 0.3519 | 13000 | 1.7508        | 3.7326     | 0.7006        | 0.7247          | 0.4583          |
+| 0.3546 | 13100 | 1.7667        | -          | -             | -               | -               |
+| 0.3573 | 13200 | 1.6415        | -          | -             | -               | -               |
+| 0.3600 | 13300 | 1.7501        | -          | -             | -               | -               |
+| 0.3627 | 13400 | 1.8451        | -          | -             | -               | -               |
+| 0.3654 | 13500 | 1.7146        | -          | -             | -               | -               |
+| 0.3681 | 13600 | 1.6837        | -          | -             | -               | -               |
+| 0.3708 | 13700 | 1.92          | -          | -             | -               | -               |
+| 0.3735 | 13800 | 1.6925        | -          | -             | -               | -               |
+| 0.3763 | 13900 | 1.7799        | -          | -             | -               | -               |
+| 0.3790 | 14000 | 1.527         | 3.7260     | 0.6989        | 0.727           | 0.4510          |
+| 0.3817 | 14100 | 1.7222        | -          | -             | -               | -               |
+| 0.3844 | 14200 | 1.8278        | -          | -             | -               | -               |
+| 0.3871 | 14300 | 1.7669        | -          | -             | -               | -               |
+| 0.3898 | 14400 | 1.5856        | -          | -             | -               | -               |
+| 0.3925 | 14500 | 1.8234        | -          | -             | -               | -               |
+| 0.3952 | 14600 | 1.7151        | -          | -             | -               | -               |
+| 0.3979 | 14700 | 1.6432        | -          | -             | -               | -               |
+| 0.4006 | 14800 | 1.9005        | -          | -             | -               | -               |
+| 0.4033 | 14900 | 1.6946        | -          | -             | -               | -               |
+| 0.4060 | 15000 | 1.5543        | 3.7222     | 0.6969        | 0.7275          | 0.4634          |
+| 0.4087 | 15100 | 1.6736        | -          | -             | -               | -               |
+| 0.4114 | 15200 | 1.8898        | -          | -             | -               | -               |
+| 0.4141 | 15300 | 1.7224        | -          | -             | -               | -               |
+| 0.4169 | 15400 | 1.7909        | -          | -             | -               | -               |
+| 0.4196 | 15500 | 1.6555        | -          | -             | -               | -               |
+| 0.4223 | 15600 | 1.523         | -          | -             | -               | -               |
+| 0.4250 | 15700 | 1.7539        | -          | -             | -               | -               |
+| 0.4277 | 15800 | 1.5763        | -          | -             | -               | -               |
+| 0.4304 | 15900 | 1.7247        | -          | -             | -               | -               |
+| 0.4331 | 16000 | 1.876         | 3.7105     | 0.6977        | 0.7263          | 0.4636          |
+| 0.4358 | 16100 | 1.772         | -          | -             | -               | -               |
+| 0.4385 | 16200 | 1.6774        | -          | -             | -               | -               |
+| 0.4412 | 16300 | 1.7602        | -          | -             | -               | -               |
+| 0.4439 | 16400 | 1.705         | -          | -             | -               | -               |
+| 0.4466 | 16500 | 1.7893        | -          | -             | -               | -               |
+| 0.4493 | 16600 | 1.653         | -          | -             | -               | -               |
+| 0.4520 | 16700 | 1.8326        | -          | -             | -               | -               |
+| 0.4547 | 16800 | 1.5326        | -          | -             | -               | -               |
+| 0.4575 | 16900 | 1.8251        | -          | -             | -               | -               |
+| 0.4602 | 17000 | 1.766         | 3.7193     | 0.6973        | 0.7257          | 0.4655          |
+| 0.4629 | 17100 | 1.7162        | -          | -             | -               | -               |
+| 0.4656 | 17200 | 1.6969        | -          | -             | -               | -               |
+| 0.4683 | 17300 | 1.5172        | -          | -             | -               | -               |
+| 0.4710 | 17400 | 1.7102        | -          | -             | -               | -               |
+| 0.4737 | 17500 | 1.8369        | -          | -             | -               | -               |
+| 0.4764 | 17600 | 1.8069        | -          | -             | -               | -               |
+| 0.4791 | 17700 | 1.6299        | -          | -             | -               | -               |
+| 0.4818 | 17800 | 1.8474        | -          | -             | -               | -               |
+| 0.4845 | 17900 | 1.5864        | -          | -             | -               | -               |
+| 0.4872 | 18000 | 1.7455        | 3.7087     | 0.6986        | 0.7249          | 0.4626          |
+| 0.4899 | 18100 | 1.8263        | -          | -             | -               | -               |
+| 0.4926 | 18200 | 1.8548        | -          | -             | -               | -               |
+| 0.4954 | 18300 | 1.6442        | -          | -             | -               | -               |
+| 0.4981 | 18400 | 1.7467        | -          | -             | -               | -               |
+| 0.5008 | 18500 | 1.6174        | -          | -             | -               | -               |
+| 0.5035 | 18600 | 1.4465        | -          | -             | -               | -               |
+| 0.5062 | 18700 | 1.8866        | -          | -             | -               | -               |
+| 0.5089 | 18800 | 1.72          | -          | -             | -               | -               |
+| 0.5116 | 18900 | 1.7466        | -          | -             | -               | -               |
+| 0.5143 | 19000 | 1.9124        | 3.7247     | 0.6979        | 0.725           | 0.4602          |
+| 0.5170 | 19100 | 1.5687        | -          | -             | -               | -               |
+| 0.5197 | 19200 | 1.6391        | -          | -             | -               | -               |
+| 0.5224 | 19300 | 1.8248        | -          | -             | -               | -               |
+| 0.5251 | 19400 | 1.6231        | -          | -             | -               | -               |
+| 0.5278 | 19500 | 1.6152        | -          | -             | -               | -               |
+| 0.5305 | 19600 | 1.639         | -          | -             | -               | -               |
+| 0.5332 | 19700 | 1.6098        | -          | -             | -               | -               |
+| 0.5360 | 19800 | 1.6619        | -          | -             | -               | -               |
+| 0.5387 | 19900 | 1.6997        | -          | -             | -               | -               |
+| 0.5414 | 20000 | 1.718         | 3.7259     | 0.6989        | 0.7264          | 0.4660          |
+| 0.5441 | 20100 | 1.634         | -          | -             | -               | -               |
+| 0.5468 | 20200 | 1.7865        | -          | -             | -               | -               |
+| 0.5495 | 20300 | 1.8573        | -          | -             | -               | -               |
+| 0.5522 | 20400 | 1.5575        | -          | -             | -               | -               |
+| 0.5549 | 20500 | 1.6594        | -          | -             | -               | -               |
+| 0.5576 | 20600 | 1.8793        | -          | -             | -               | -               |
+| 0.5603 | 20700 | 1.7643        | -          | -             | -               | -               |
+| 0.5630 | 20800 | 1.538         | -          | -             | -               | -               |
+| 0.5657 | 20900 | 1.8634        | -          | -             | -               | -               |
+| 0.5684 | 21000 | 1.916         | 3.7223     | 0.6982        | 0.7258          | 0.4650          |
+| 0.5711 | 21100 | 1.5947        | -          | -             | -               | -               |
+| 0.5738 | 21200 | 1.5321        | -          | -             | -               | -               |
+| 0.5766 | 21300 | 1.7004        | -          | -             | -               | -               |
+| 0.5793 | 21400 | 1.6947        | -          | -             | -               | -               |
+| 0.5820 | 21500 | 1.5228        | -          | -             | -               | -               |
+| 0.5847 | 21600 | 1.7152        | -          | -             | -               | -               |
+| 0.5874 | 21700 | 1.6883        | -          | -             | -               | -               |
+| 0.5901 | 21800 | 1.6779        | -          | -             | -               | -               |
+| 0.5928 | 21900 | 1.7323        | -          | -             | -               | -               |
+| 0.5955 | 22000 | 1.9633        | 3.7266     | 0.6996        | 0.7288          | 0.4635          |
+| 0.5982 | 22100 | 1.7498        | -          | -             | -               | -               |
+| 0.6009 | 22200 | 1.7513        | -          | -             | -               | -               |
+| 0.6036 | 22300 | 1.7078        | -          | -             | -               | -               |
+| 0.6063 | 22400 | 1.6438        | -          | -             | -               | -               |
+| 0.6090 | 22500 | 1.6743        | -          | -             | -               | -               |
+| 0.6117 | 22600 | 1.6701        | -          | -             | -               | -               |
+| 0.6145 | 22700 | 1.7871        | -          | -             | -               | -               |
+| 0.6172 | 22800 | 1.6247        | -          | -             | -               | -               |
+| 0.6199 | 22900 | 1.7817        | -          | -             | -               | -               |
+| 0.6226 | 23000 | 1.6606        | 3.7321     | 0.6993        | 0.7286          | 0.4614          |
+| 0.6253 | 23100 | 1.8987        | -          | -             | -               | -               |
+| 0.6280 | 23200 | 1.6494        | -          | -             | -               | -               |
+| 0.6307 | 23300 | 1.6776        | -          | -             | -               | -               |
+| 0.6334 | 23400 | 1.75          | -          | -             | -               | -               |
+| 0.6361 | 23500 | 1.5131        | -          | -             | -               | -               |
+| 0.6388 | 23600 | 1.7946        | -          | -             | -               | -               |
+| 0.6415 | 23700 | 1.665         | -          | -             | -               | -               |
+| 0.6442 | 23800 | 1.6681        | -          | -             | -               | -               |
+| 0.6469 | 23900 | 1.8255        | -          | -             | -               | -               |
+| 0.6496 | 24000 | 1.6759        | 3.7227     | 0.7017        | 0.7281          | 0.4625          |
+| 0.6523 | 24100 | 1.554         | -          | -             | -               | -               |
+| 0.6551 | 24200 | 1.6435        | -          | -             | -               | -               |
+| 0.6578 | 24300 | 1.8224        | -          | -             | -               | -               |
+| 0.6605 | 24400 | 1.6186        | -          | -             | -               | -               |
+| 0.6632 | 24500 | 1.7156        | -          | -             | -               | -               |
+| 0.6659 | 24600 | 1.5247        | -          | -             | -               | -               |
+| 0.6686 | 24700 | 1.6264        | -          | -             | -               | -               |
+| 0.6713 | 24800 | 1.7673        | -          | -             | -               | -               |
+| 0.6740 | 24900 | 1.8072        | -          | -             | -               | -               |
+| 0.6767 | 25000 | 1.765         | 3.7407     | 0.7026        | 0.7283          | 0.4589          |
+| 0.6794 | 25100 | 1.6422        | -          | -             | -               | -               |
+| 0.6821 | 25200 | 1.7846        | -          | -             | -               | -               |
+| 0.6848 | 25300 | 1.7366        | -          | -             | -               | -               |
+| 0.6875 | 25400 | 1.7839        | -          | -             | -               | -               |
+| 0.6902 | 25500 | 1.441         | -          | -             | -               | -               |
+| 0.6930 | 25600 | 1.5533        | -          | -             | -               | -               |
+| 0.6957 | 25700 | 1.6922        | -          | -             | -               | -               |
+| 0.6984 | 25800 | 1.5544        | -          | -             | -               | -               |
+| 0.7011 | 25900 | 1.456         | -          | -             | -               | -               |
+| 0.7038 | 26000 | 1.6494        | 3.7274     | 0.7059        | 0.7268          | 0.4661          |
+| 0.7065 | 26100 | 1.6963        | -          | -             | -               | -               |
+| 0.7092 | 26200 | 1.7892        | -          | -             | -               | -               |
+| 0.7119 | 26300 | 1.6669        | -          | -             | -               | -               |
+| 0.7146 | 26400 | 1.6758        | -          | -             | -               | -               |
+| 0.7173 | 26500 | 1.6322        | -          | -             | -               | -               |
+| 0.7200 | 26600 | 1.5416        | -          | -             | -               | -               |
+| 0.7227 | 26700 | 1.681         | -          | -             | -               | -               |
+| 0.7254 | 26800 | 1.5159        | -          | -             | -               | -               |
+| 0.7281 | 26900 | 1.715         | -          | -             | -               | -               |
+| 0.7308 | 27000 | 1.6164        | 3.7456     | 0.7061        | 0.7257          | 0.4570          |
+| 0.7336 | 27100 | 1.6784        | -          | -             | -               | -               |
+| 0.7363 | 27200 | 1.5886        | -          | -             | -               | -               |
+| 0.7390 | 27300 | 1.6736        | -          | -             | -               | -               |
+| 0.7417 | 27400 | 1.5659        | -          | -             | -               | -               |
+| 0.7444 | 27500 | 1.6552        | -          | -             | -               | -               |
+| 0.7471 | 27600 | 1.5672        | -          | -             | -               | -               |
+| 0.7498 | 27700 | 1.5873        | -          | -             | -               | -               |
+| 0.7525 | 27800 | 1.6746        | -          | -             | -               | -               |
+| 0.7552 | 27900 | 1.7503        | -          | -             | -               | -               |
+| 0.7579 | 28000 | 1.7287        | 3.7390     | 0.7076        | 0.7244          | 0.4636          |
+| 0.7606 | 28100 | 1.6216        | -          | -             | -               | -               |
+| 0.7633 | 28200 | 1.6101        | -          | -             | -               | -               |
+| 0.7660 | 28300 | 1.5651        | -          | -             | -               | -               |
+| 0.7687 | 28400 | 1.5659        | -          | -             | -               | -               |
+| 0.7714 | 28500 | 1.5248        | -          | -             | -               | -               |
+| 0.7742 | 28600 | 1.3725        | -          | -             | -               | -               |
+| 0.7769 | 28700 | 1.7881        | -          | -             | -               | -               |
+| 0.7796 | 28800 | 1.739         | -          | -             | -               | -               |
+| 0.7823 | 28900 | 1.6464        | -          | -             | -               | -               |
+| 0.7850 | 29000 | 1.6841        | 3.7212     | 0.7073        | 0.7247          | 0.4626          |
+| 0.7877 | 29100 | 1.6254        | -          | -             | -               | -               |
+| 0.7904 | 29200 | 1.6728        | -          | -             | -               | -               |
+| 0.7931 | 29300 | 1.5605        | -          | -             | -               | -               |
+| 0.7958 | 29400 | 1.687         | -          | -             | -               | -               |
+| 0.7985 | 29500 | 1.7799        | -          | -             | -               | -               |
+| 0.8012 | 29600 | 1.6792        | -          | -             | -               | -               |
+| 0.8039 | 29700 | 1.5241        | -          | -             | -               | -               |
+| 0.8066 | 29800 | 1.6341        | -          | -             | -               | -               |
+| 0.8093 | 29900 | 1.5571        | -          | -             | -               | -               |
+| 0.8121 | 30000 | 1.5228        | 3.7397     | 0.7105        | 0.7234          | 0.4682          |
+| 0.8148 | 30100 | 1.5988        | -          | -             | -               | -               |
+| 0.8175 | 30200 | 1.4222        | -          | -             | -               | -               |
+| 0.8202 | 30300 | 1.4629        | -          | -             | -               | -               |
+| 0.8229 | 30400 | 1.6381        | -          | -             | -               | -               |
+| 0.8256 | 30500 | 1.4585        | -          | -             | -               | -               |
+| 0.8283 | 30600 | 1.6774        | -          | -             | -               | -               |
+| 0.8310 | 30700 | 1.811         | -          | -             | -               | -               |
+| 0.8337 | 30800 | 1.5872        | -          | -             | -               | -               |
+| 0.8364 | 30900 | 1.4762        | -          | -             | -               | -               |
+| 0.8391 | 31000 | 1.7079        | 3.7256     | 0.7128        | 0.7215          | 0.4645          |
+| 0.8418 | 31100 | 1.4948        | -          | -             | -               | -               |
+| 0.8445 | 31200 | 1.4556        | -          | -             | -               | -               |
+| 0.8472 | 31300 | 1.5191        | -          | -             | -               | -               |
+| 0.8499 | 31400 | 1.598         | -          | -             | -               | -               |
+| 0.8527 | 31500 | 1.6586        | -          | -             | -               | -               |
+| 0.8554 | 31600 | 1.6893        | -          | -             | -               | -               |
+| 0.8581 | 31700 | 1.7764        | -          | -             | -               | -               |
+| 0.8608 | 31800 | 1.3632        | -          | -             | -               | -               |
+| 0.8635 | 31900 | 1.6681        | -          | -             | -               | -               |
+| 0.8662 | 32000 | 1.6232        | 3.7358     | 0.7161        | 0.7232          | 0.4651          |
+| 0.8689 | 32100 | 1.4556        | -          | -             | -               | -               |
+| 0.8716 | 32200 | 1.8698        | -          | -             | -               | -               |
+| 0.8743 | 32300 | 1.7566        | -          | -             | -               | -               |
+| 0.8770 | 32400 | 1.6082        | -          | -             | -               | -               |
+| 0.8797 | 32500 | 1.6465        | -          | -             | -               | -               |
+| 0.8824 | 32600 | 1.5018        | -          | -             | -               | -               |
+| 0.8851 | 32700 | 1.8482        | -          | -             | -               | -               |
+| 0.8878 | 32800 | 1.5147        | -          | -             | -               | -               |
+| 0.8905 | 32900 | 1.699         | -          | -             | -               | -               |
+| 0.8933 | 33000 | 1.5738        | 3.7323     | 0.7176        | 0.7246          | 0.4657          |
+| 0.8960 | 33100 | 1.635         | -          | -             | -               | -               |
+| 0.8987 | 33200 | 1.7069        | -          | -             | -               | -               |
+| 0.9014 | 33300 | 1.6272        | -          | -             | -               | -               |
+| 0.9041 | 33400 | 1.7648        | -          | -             | -               | -               |
+| 0.9068 | 33500 | 1.6683        | -          | -             | -               | -               |
+| 0.9095 | 33600 | 1.4867        | -          | -             | -               | -               |
+| 0.9122 | 33700 | 1.6677        | -          | -             | -               | -               |
+| 0.9149 | 33800 | 1.5527        | -          | -             | -               | -               |
+| 0.9176 | 33900 | 1.6804        | -          | -             | -               | -               |
+| 0.9203 | 34000 | 1.425         | 3.7477     | 0.7172        | 0.7231          | 0.4596          |
+| 0.9230 | 34100 | 1.771         | -          | -             | -               | -               |
+| 0.9257 | 34200 | 1.5767        | -          | -             | -               | -               |
+| 0.9284 | 34300 | 1.5424        | -          | -             | -               | -               |
+| 0.9312 | 34400 | 1.5985        | -          | -             | -               | -               |
+| 0.9339 | 34500 | 1.6763        | -          | -             | -               | -               |
+| 0.9366 | 34600 | 1.6608        | -          | -             | -               | -               |
+| 0.9393 | 34700 | 1.7736        | -          | -             | -               | -               |
+| 0.9420 | 34800 | 1.8955        | -          | -             | -               | -               |
+| 0.9447 | 34900 | 1.5688        | -          | -             | -               | -               |
+| 0.9474 | 35000 | 1.6123        | 3.7410     | 0.7196        | 0.7226          | 0.4671          |
+| 0.9501 | 35100 | 1.7264        | -          | -             | -               | -               |
+| 0.9528 | 35200 | 1.5511        | -          | -             | -               | -               |
+| 0.9555 | 35300 | 1.6409        | -          | -             | -               | -               |
+| 0.9582 | 35400 | 1.47          | -          | -             | -               | -               |
+| 0.9609 | 35500 | 1.8675        | -          | -             | -               | -               |
+| 0.9636 | 35600 | 1.6868        | -          | -             | -               | -               |
+| 0.9663 | 35700 | 1.744         | -          | -             | -               | -               |
+| 0.9690 | 35800 | 1.6734        | -          | -             | -               | -               |
+| 0.9718 | 35900 | 1.4154        | -          | -             | -               | -               |
+| 0.9745 | 36000 | 1.4793        | 3.7393     | 0.7190        | 0.7223          | 0.4677          |
+| 0.9772 | 36100 | 1.7126        | -          | -             | -               | -               |
+| 0.9799 | 36200 | 1.7037        | -          | -             | -               | -               |
+| 0.9826 | 36300 | 1.6306        | -          | -             | -               | -               |
+| 0.9853 | 36400 | 1.7783        | -          | -             | -               | -               |
+| 0.9880 | 36500 | 1.5751        | -          | -             | -               | -               |
+| 0.9907 | 36600 | 1.6079        | -          | -             | -               | -               |
+| 0.9934 | 36700 | 1.7162        | -          | -             | -               | -               |
+| 0.9961 | 36800 | 1.447         | -          | -             | -               | -               |
+| 0.9988 | 36900 | 1.6155        | -          | -             | -               | -               |
+| 1.0015 | 37000 | 1.7294        | 3.7512     | 0.7177        | 0.7236          | 0.4659          |
+
+</details>
+
+### Framework Versions
+- Python: 3.10.12
+- Sentence Transformers: 3.0.1
+- Transformers: 4.38.2
+- PyTorch: 2.1.2+cu121
+- Accelerate: 0.27.2
+- Datasets: 2.19.1
+- Tokenizers: 0.15.2
+
+## Citation
+
+### BibTeX
+
+#### Sentence Transformers
+```bibtex
+@inproceedings{reimers-2019-sentence-bert,
+    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
+    author = "Reimers, Nils and Gurevych, Iryna",
+    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
+    month = "11",
+    year = "2019",
+    publisher = "Association for Computational Linguistics",
+    url = "https://arxiv.org/abs/1908.10084",
+}
+```
+
+#### CachedMultipleNegativesRankingLoss
+```bibtex
+@misc{gao2021scaling,
+    title={Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup}, 
+    author={Luyu Gao and Yunyi Zhang and Jiawei Han and Jamie Callan},
+    year={2021},
+    eprint={2101.06983},
+    archivePrefix={arXiv},
+    primaryClass={cs.LG}
+}
+```
+
+#### AnglELoss
+```bibtex
+@misc{li2023angleoptimized,
+    title={AnglE-optimized Text Embeddings}, 
+    author={Xianming Li and Jing Li},
+    year={2023},
+    eprint={2309.12871},
+    archivePrefix={arXiv},
+    primaryClass={cs.CL}
+}
+```
+
+<!--
+## Glossary
+
+*Clearly define terms in order to be accessible across audiences.*
+-->
+
+<!--
+## Model Card Authors
+
+*Lists the people who create the model card, providing recognition and accountability for the detailed work that goes into its construction.*
+-->
+
+<!--
+## Model Card Contact
+
+*Provides a way for people who have updates to the Model Card, suggestions, or questions, to contact the Model Card authors.*
+-->

config.json

@@ -0,0 +1,58 @@
+{
+  "_name_or_path": "models/nomic-embed-text-esci/checkpoint-37000",
+  "activation_function": "swiglu",
+  "architectures": [
+    "NomicBertModel"
+  ],
+  "attn_pdrop": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_hf_nomic_bert.NomicBertConfig",
+    "AutoModel": "modeling_hf_nomic_bert.NomicBertModel",
+    "AutoModelForMaskedLM": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForPreTraining"
+  },
+  "bos_token_id": null,
+  "causal": false,
+  "dense_seq_output": true,
+  "embd_pdrop": 0.0,
+  "eos_token_id": null,
+  "fused_bias_fc": true,
+  "fused_dropout_add_ln": true,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-12,
+  "max_trained_positions": 2048,
+  "mlp_fc1_bias": false,
+  "mlp_fc2_bias": false,
+  "model_type": "nomic_bert",
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": 3072,
+  "n_layer": 12,
+  "n_positions": 8192,
+  "pad_vocab_size_multiple": 64,
+  "parallel_block": false,
+  "parallel_block_tied_norm": false,
+  "prenorm": false,
+  "qkv_proj_bias": false,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.0,
+  "rotary_emb_base": 1000,
+  "rotary_emb_fraction": 1.0,
+  "rotary_emb_interleaved": false,
+  "rotary_emb_scale_base": null,
+  "rotary_scaling_factor": null,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.0,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "use_flash_attn": true,
+  "use_rms_norm": false,
+  "use_xentropy": true,
+  "vocab_size": 30528
+}

config_sentence_transformers.json

@@ -0,0 +1,10 @@
+{
+  "__version__": {
+    "sentence_transformers": "3.0.1",
+    "transformers": "4.38.2",
+    "pytorch": "2.1.2+cu121"
+  },
+  "prompts": {},
+  "default_prompt_name": null,
+  "similarity_fn_name": null
+}

configuration_hf_nomic_bert.py

@@ -0,0 +1,56 @@
+from transformers import GPT2Config
+
+
+class NomicBertConfig(GPT2Config):
+    model_type = "nomic_bert"
+
+    def __init__(
+        self,
+        prenorm=False,
+        parallel_block=False,
+        parallel_block_tied_norm=False,
+        rotary_emb_fraction=0.0,
+        fused_dropout_add_ln=False,
+        fused_bias_fc=False,
+        use_flash_attn=False,
+        use_xentropy=False,
+        qkv_proj_bias=True,
+        rotary_emb_base=10_000,
+        rotary_emb_scale_base=None,
+        rotary_emb_interleaved=False,
+        mlp_fc1_bias=True,
+        mlp_fc2_bias=True,
+        use_rms_norm=False,
+        causal=False,
+        type_vocab_size=2,
+        dense_seq_output=True,
+        pad_vocab_size_multiple=1,
+        tie_word_embeddings=True,
+        rotary_scaling_factor=None,
+        max_trained_positions=2048,
+        **kwargs,
+    ):
+        self.prenorm = prenorm
+        self.parallel_block = parallel_block
+        self.parallel_block_tied_norm = parallel_block_tied_norm
+        self.rotary_emb_fraction = rotary_emb_fraction
+        self.tie_word_embeddings = tie_word_embeddings
+        self.fused_dropout_add_ln = fused_dropout_add_ln
+        self.fused_bias_fc = fused_bias_fc
+        self.use_flash_attn = use_flash_attn
+        self.use_xentropy = use_xentropy
+        self.qkv_proj_bias = qkv_proj_bias
+        self.rotary_emb_base = rotary_emb_base
+        self.rotary_emb_scale_base = rotary_emb_scale_base
+        self.rotary_emb_interleaved = rotary_emb_interleaved
+        self.mlp_fc1_bias = mlp_fc1_bias
+        self.mlp_fc2_bias = mlp_fc2_bias
+        self.use_rms_norm = use_rms_norm
+        self.causal = causal
+        self.type_vocab_size = type_vocab_size
+        self.dense_seq_output = dense_seq_output
+        self.pad_vocab_size_multiple = pad_vocab_size_multiple
+        self.rotary_scaling_factor = rotary_scaling_factor
+        self.max_trained_positions = max_trained_positions
+
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a3ba911558d833fc7ede5bbf4cccc67f19a69e1f7b5495aab145106c839b8aa
+size 546938168

modeling_hf_nomic_bert.py

@@ -0,0 +1,2071 @@
+# Copyright (c) 2022, Tri Dao.
+# This BERT implementation is based on our MLPerf 2.0 and MLPerf 2.1 BERT implementation.
+# https://github.com/mlcommons/training_results_v2.0/blob/main/HazyResearch/benchmarks/bert/implementations/pytorch/modeling.py
+# https://github.com/mlcommons/training_results_v2.1/blob/main/Azure-HazyResearch/benchmarks/bert/implementations/ND96amsr_A100_v4/modeling.py
+
+import logging
+
+# Inspired by https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py
+import math
+import numpy as np
+import collections
+import os
+import re
+from collections import OrderedDict
+from functools import partial
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from safetensors.torch import load_file as safe_load_file
+from transformers import GPT2Config, PreTrainedModel, ViTModel, ViTConfig
+from transformers.models.bert.modeling_bert import (
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    MaskedLMOutput,
+    SequenceClassifierOutput,
+)
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME
+from transformers.utils.hub import cached_file, get_checkpoint_shard_files
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from torch.nn.modules.utils import _pair
+
+from .configuration_hf_nomic_bert import NomicBertConfig
+
+logger = logging.getLogger(__name__)
+
+
+# adapted from flash attention, added safe serialization option for hf models
+def state_dict_from_pretrained(model_name, safe_serialization=False, device=None, dtype=None):
+    # If not fp32, then we don't want to load directly to the GPU
+    mapped_device = "cpu" if dtype not in [torch.float32, None] else device
+    is_sharded = False
+    load_safe = False
+    resolved_archive_file = None
+
+    weights_path = os.path.join(model_name, WEIGHTS_NAME)
+    weights_index_path = os.path.join(model_name, WEIGHTS_INDEX_NAME)
+    safe_weights_path = os.path.join(model_name, SAFE_WEIGHTS_NAME)
+    safe_weights_index_path = os.path.join(model_name, SAFE_WEIGHTS_INDEX_NAME)
+
+    if os.path.isfile(weights_path):
+        resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
+    elif os.path.isfile(weights_index_path):
+        resolved_archive_file = cached_file(model_name, WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False)
+        is_sharded = True
+    elif os.path.isfile(safe_weights_path):
+        resolved_archive_file = cached_file(model_name, SAFE_WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
+        load_safe = True
+    elif os.path.isfile(safe_weights_index_path):
+        resolved_archive_file = cached_file(
+            model_name, SAFE_WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False
+        )
+        is_sharded = True
+        load_safe = True
+    else:  # Try loading from HF hub instead of from local files
+        resolved_archive_file = None
+        for weight_name in [WEIGHTS_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]:
+            resolved_archive_file = cached_file(
+                model_name, weight_name, _raise_exceptions_for_missing_entries=False
+            )
+            if resolved_archive_file is not None:
+                if weight_name in [SAFE_WEIGHTS_NAME, SAFE_WEIGHTS_INDEX_NAME]:
+                    load_safe = True
+                if weight_name in [WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]:
+                    is_sharded = True
+                break
+
+    if resolved_archive_file is None:
+        raise EnvironmentError(f"Model name {model_name} was not found.")
+
+    if load_safe:
+        loader = partial(safe_load_file, device=mapped_device)
+    else:
+        loader = partial(torch.load, map_location=mapped_device)
+
+    if is_sharded:
+        # resolved_archive_file becomes a list of files that point to the different
+        # checkpoint shards in this case.
+        resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(model_name, resolved_archive_file)
+        state_dict = {}
+        for sharded_file in resolved_archive_file:
+            state_dict.update(loader(sharded_file))
+    else:
+        state_dict = loader(resolved_archive_file)
+    # Convert dtype before moving to GPU to save memory
+    if dtype is not None:
+        state_dict = {k: v.to(dtype=dtype) for k, v in state_dict.items()}
+    state_dict = {k: v.to(device=device) for k, v in state_dict.items()}
+    return state_dict
+
+
+def filter_shapes(state_dict, model):
+    """
+    Filters the state dict to match the current model shape.
+    """
+    filtered_state_dict = {}
+    for key, value in state_dict.items():
+        if key in model.state_dict():
+            if value.shape == model.state_dict()[key].shape:
+                filtered_state_dict[key] = value
+    return filtered_state_dict
+
+
+def remap_bert_state_dict(
+    state_dict,
+    config,
+    remove_bert=False,
+    remove_cls_weights=False,
+    add_pooling_layer=False,
+):
+    """
+    Map the state_dict of a Huggingface BERT model to be flash_attn compatible.
+    """
+
+    def add_bert_prefix(key):
+        # prepend bert. to the key
+        if key.startswith("bert.") or key.startswith("cls."):
+            return key
+        return f"bert.{key}"
+
+    state_dict = OrderedDict((add_bert_prefix(k), v) for k, v in state_dict.items())
+
+    # LayerNorm
+    def key_mapping_ln_gamma_beta(key):
+        key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key)
+        key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key)
+        return key
+
+    state_dict = OrderedDict((key_mapping_ln_gamma_beta(k), v) for k, v in state_dict.items())
+
+    # Layers
+    def key_mapping_layers(key):
+        return re.sub(r"^bert.encoder.layer\.", "bert.encoder.layers.", key)
+
+    state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items())
+
+    # LayerNorm
+    def key_mapping_ln(key):
+        key = re.sub(r"^bert.embeddings.LayerNorm.", "bert.emb_ln.", key)
+        key = re.sub(
+            r"^bert.encoder.layers.(\d+).attention.output.LayerNorm.(weight|bias)",
+            r"bert.encoder.layers.\1.norm1.\2",
+            key,
+        )
+        key = re.sub(
+            r"^bert.encoder.layers.(\d+).output.LayerNorm.(weight|bias)",
+            r"bert.encoder.layers.\1.norm2.\2",
+            key,
+        )
+        key = re.sub(
+            r"^cls.predictions.transform.LayerNorm.(weight|bias)",
+            r"cls.predictions.transform.layer_norm.\1",
+            key,
+        )
+        return key
+
+    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
+
+    # MLP
+    def key_mapping_mlp(key):
+        key = re.sub(
+            r"^bert.encoder.layers.(\d+).intermediate.dense.(weight|bias)",
+            r"bert.encoder.layers.\1.mlp.fc1.\2",
+            key,
+        )
+        key = re.sub(
+            r"^bert.encoder.layers.(\d+).output.dense.(weight|bias)",
+            r"bert.encoder.layers.\1.mlp.fc2.\2",
+            key,
+        )
+        return key
+
+    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())
+
+    # Attention
+    last_layer_subset = getattr(config, "last_layer_subset", False)
+    for d in range(config.num_hidden_layers):
+        if f"bert.encoder.layers.{d}.attention.self.query.weight" not in state_dict:
+            continue
+        Wq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.weight")
+        Wk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.weight")
+        Wv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.weight")
+        bq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.bias")
+        bk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.bias")
+        bv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.bias")
+        if not (last_layer_subset and d == config.num_hidden_layers - 1):
+            state_dict[f"bert.encoder.layers.{d}.attn.Wqkv.weight"] = torch.cat([Wq, Wk, Wv], dim=0)
+            state_dict[f"bert.encoder.layers.{d}.attn.Wqkv.bias"] = torch.cat([bq, bk, bv], dim=0)
+        else:
+            state_dict[f"bert.encoder.layers.{d}.attn.Wq.weight"] = Wq
+            state_dict[f"bert.encoder.layers.{d}.attn.Wkv.weight"] = torch.cat([Wk, Wv], dim=0)
+            state_dict[f"bert.encoder.layers.{d}.attn.Wq.bias"] = bq
+            state_dict[f"bert.encoder.layers.{d}.attn.Wkv.bias"] = torch.cat([bk, bv], dim=0)
+
+    def key_mapping_attn(key):
+        return re.sub(
+            r"^bert.encoder.layers.(\d+).attention.output.dense.(weight|bias)",
+            r"bert.encoder.layers.\1.attn.out_proj.\2",
+            key,
+        )
+
+    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())
+
+    def key_mapping_decoder_bias(key):
+        return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key)
+
+    # remove nsp weights, we don't use
+    state_dict.pop("cls.seq_relationship.weight", None)
+    state_dict.pop("cls.seq_relationship.bias", None)
+    state_dict.pop("bert.embeddings.position_ids", None)
+
+    state_dict = OrderedDict((key_mapping_decoder_bias(k), v) for k, v in state_dict.items())
+
+    if remove_cls_weights:
+        cls_weights = [
+            "cls.predictions.decoder.bias",
+            "cls.predictions.transform.dense.weight",
+            "cls.predictions.transform.dense.bias",
+            "cls.predictions.transform.layer_norm.weight",
+            "cls.predictions.transform.layer_norm.bias",
+            "cls.predictions.decoder.weight",
+        ]
+        for weight in cls_weights:
+            state_dict.pop(weight, None)
+
+    # Word embedding
+    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
+    if pad_vocab_size_multiple > 1:
+        word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
+        state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(
+            word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])
+        )
+        if not remove_cls_weights:
+            decoder_weight = state_dict["cls.predictions.decoder.weight"]
+            state_dict["cls.predictions.decoder.weight"] = F.pad(
+                decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])
+            )
+            # If the vocab was padded, we want to set the decoder bias for those padded indices to be
+            # strongly negative (i.e. the decoder shouldn't predict those indices).
+            # TD [2022-05-09]: I don't think it affects the MLPerf training.
+            if "cls.predictions.decoder.bias" in state_dict:
+                decoder_bias = state_dict["cls.predictions.decoder.bias"]
+                state_dict["cls.predictions.decoder.bias"] = F.pad(
+                    decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
+                )
+
+    if add_pooling_layer is False:
+        pooler_weights = [
+            "bert.pooler.dense.weight",
+            "bert.pooler.dense.bias",
+        ]
+        for key in pooler_weights:
+            state_dict.pop(key, None)
+
+    if remove_bert:
+
+        def remove_bert_prefix(key):
+            key = re.sub(r"^bert.", "", key)
+            return key
+
+        state_dict = OrderedDict((remove_bert_prefix(k), v) for k, v in state_dict.items())
+
+    return state_dict
+
+    
+def _trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        print("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    # Values are generated by using a truncated uniform distribution and
+    # then using the inverse CDF for the normal distribution.
+    # Get upper and lower cdf values
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated
+    # standard normal
+    tensor.erfinv_()
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    tensor.clamp_(min=a, max=b)
+    return tensor
+
+def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.):
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+
+    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
+    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
+    and the result is subsquently scaled and shifted by the mean and std args.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    with torch.no_grad():
+        _trunc_normal_(tensor, 0, 1.0, a, b)
+        tensor.mul_(std).add_(mean)
+    return tensor
+
+
+class NomicBertPreTrainedModel(PreTrainedModel):
+    """An abstract class to handle weights initialization and
+    a simple interface for dowloading and loading pretrained models.
+    """
+
+    config_class = NomicBertConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config)
+        if not isinstance(config, GPT2Config):
+            raise ValueError(
+                "Parameter config in `{}(config)` should be an instance of class `GPT2Config`. "
+                "To create a model from a Google pretrained model use "
+                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    self.__class__.__name__, self.__class__.__name__
+                )
+            )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, model_name, config=None, *inputs, **kwargs):
+        """
+        Instantiate a NomicBertPreTrainedModel from a pre-trained model file or a pytorch state dict.
+        Download and cache the pre-trained model file if needed.
+
+        Params:
+            pretrained_model_name_or_path: either:
+                - a path or url to a pretrained model archive containing:
+                    . `bert_config.json` a configuration file for the model
+                    . `pytorch_model.bin` a PyTorch dump of a NomicBertForPretraining instance
+                - a path or url to a pretrained model archive containing:
+                    . `bert_config.json` a configuration file for the model
+                    . `model.chkpt` a TensorFlow checkpoint
+            *inputs, **kwargs: additional input for the specific NomicBert class
+                (ex: num_labels for NomicBertForSequenceClassification)
+        """
+        # Instantiate model.
+        if config is None:
+            config = cls.config_class.from_pretrained(model_name)
+        remove_cls = cls != NomicBertForPreTraining
+        remove_bert_prefix = cls != NomicBertForPreTraining and cls != NomicBertForSequenceClassification
+        ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
+        num_labels = kwargs.pop("num_labels", None)
+        rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
+        strict = kwargs.pop("strict", True)
+        if rotary_scaling_factor:
+            config.rotary_scaling_factor = rotary_scaling_factor
+
+        if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
+            config.n_positions = 2048
+        if num_labels:
+            config.num_labels = num_labels
+
+        if "add_pooling_layer" in kwargs:
+            model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
+        else:
+            if cls == NomicBertModel:
+                model = cls(config, *inputs, add_pooling_layer=False)
+            else:
+                model = cls(config, *inputs)
+        # TODO: fix this
+        # Assuming we know what we're doing when loading from disk
+        # Prob a bad assumption but i'm tired and want to train this asap
+        if os.path.exists(model_name):
+            model_path = f"{model_name}/pytorch_model.bin"
+            if os.path.exists(model_path):
+                state_dict = torch.load(f"{model_name}/pytorch_model.bin")
+            else:
+                model_path = f"{model_name}/model.safetensors"
+                if not os.path.exists(model_path):
+                    raise ValueError(f"Model path {model_path} not found")
+                state_dict = safe_load_file(model_path)
+
+            if ignore_mismatched_shapes:
+                state_dict = filter_shapes(state_dict, model)
+            load_return = model.load_state_dict(state_dict, strict=False)
+        else:
+            # TODO: can probably check config class and see if we need to remap from a bert model
+            state_dict = state_dict_from_pretrained(model_name)
+            state_dict = remap_bert_state_dict(
+                state_dict,
+                config,
+                remove_bert=remove_bert_prefix,
+                remove_cls_weights=remove_cls,
+                add_pooling_layer=getattr(config, "add_pooling_layer", False),
+            )
+            if ignore_mismatched_shapes:
+                state_dict = filter_shapes(state_dict, model)
+
+            load_return = model.load_state_dict(state_dict, strict=strict)
+        logger.warning(load_return)
+        return model
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, NomicBertEncoder):
+            module.gradient_checkpointing = value
+
+
+# https://github.com/huggingface/transformers/blob/7032e0203262ebb2ebf55da8d2e01f873973e835/src/transformers/models/bert/modeling_bert.py#L748
+def _init_weights(module, initializer_range=0.02):
+    if isinstance(module, nn.Linear):
+        nn.init.normal_(module.weight, std=initializer_range)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+        if module.padding_idx is not None:
+            nn.init.zeros_(module.weight[module.padding_idx])
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, add_cls_token=False):
+    """
+    Create 2D sin/cos positional embeddings.
+
+    Args:
+        embed_dim (`int`):
+            Embedding dimension.
+        grid_size (`int`):
+            The grid height and width.
+        add_cls_token (`bool`, *optional*, defaults to `False`):
+            Whether or not to add a classification (CLS) token.
+
+    Returns:
+        (`torch.FloatTensor` of shape (grid_size*grid_size, embed_dim) or (1+grid_size*grid_size, embed_dim): the
+        position embeddings (with or without classification token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if add_cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be even")
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be even")
+
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]:
+    """generate N-D grid in dimension order.
+
+    The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.
+
+    That is, the statement
+    [X1,X2,X3] = ndgrid(x1,x2,x3)
+
+    produces the same result as
+
+    [X2,X1,X3] = meshgrid(x2,x1,x3)
+
+    This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make
+    torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy').
+
+    """
+    try:
+        return torch.meshgrid(*tensors, indexing='ij')
+    except TypeError:
+        # old PyTorch < 1.10 will follow this path as it does not have indexing arg,
+        # the old behaviour of meshgrid was 'ij'
+        return torch.meshgrid(*tensors)
+
+def build_fourier_pos_embed(
+        feat_shape: List[int],
+        bands: Optional[torch.Tensor] = None,
+        num_bands: int = 64,
+        max_res: int = 224,
+        temperature: float = 10000.,
+        linear_bands: bool = False,
+        include_grid: bool = False,
+        in_pixels: bool = True,
+        ref_feat_shape: Optional[List[int]] = None,
+        dtype: torch.dtype = torch.float32,
+        device: Optional[torch.device] = None,
+) -> List[torch.Tensor]:
+    """
+
+    Args:
+        feat_shape: Feature shape for embedding.
+        bands: Pre-calculated frequency bands.
+        num_bands: Number of frequency bands (determines output dim).
+        max_res: Maximum resolution for pixel based freq.
+        temperature: Temperature for non-pixel freq.
+        linear_bands: Linear band spacing for pixel based freq.
+        include_grid: Include the spatial grid in output.
+        in_pixels: Output in pixel freq.
+        ref_feat_shape: Reference feature shape for resize / fine-tune.
+        dtype: Output dtype.
+        device: Output device.
+
+    Returns:
+
+    """
+    if bands is None:
+        if in_pixels:
+            bands = pixel_freq_bands(
+                num_bands,
+                float(max_res),
+                linear_bands=linear_bands,
+                device=device,
+            )
+        else:
+            bands = freq_bands(
+                num_bands,
+                temperature=temperature,
+                step=1,
+                device=device,
+            )
+    else:
+        if device is None:
+            device = bands.device
+        if dtype is None:
+            dtype = bands.dtype
+
+    if in_pixels:
+        t = [torch.linspace(-1., 1., steps=s, device=device, dtype=torch.float32) for s in feat_shape]
+    else:
+        t = [torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) for s in feat_shape]
+
+    if ref_feat_shape is not None:
+        # eva's scheme for resizing rope embeddings (ref shape = pretrain)
+        t = [x / f * r for x, f, r in zip(t, feat_shape, ref_feat_shape)]
+
+    grid = torch.stack(ndgrid(t), dim=-1)
+    grid = grid.unsqueeze(-1)
+    pos = grid * bands
+
+    pos_sin, pos_cos = pos.sin().to(dtype=dtype), pos.cos().to(dtype)
+    out = [grid, pos_sin, pos_cos] if include_grid else [pos_sin, pos_cos]
+    return out
+
+
+def build_rotary_pos_embed(
+        feat_shape: List[int],
+        bands: Optional[torch.Tensor] = None,
+        dim: int = 64,
+        max_res: int = 224,
+        temperature: float = 10000.,
+        linear_bands: bool = False,
+        in_pixels: bool = True,
+        ref_feat_shape: Optional[List[int]] = None,
+        dtype: torch.dtype = torch.float32,
+        device: Optional[torch.device] = None,
+):
+    """
+
+    Args:
+        feat_shape: Spatial shape of the target tensor for embedding.
+        bands: Optional pre-generated frequency bands
+        dim: Output dimension of embedding tensor.
+        max_res: Maximum resolution for pixel mode.
+        temperature: Temperature (inv freq) for non-pixel mode
+        linear_bands: Linearly (instead of log) spaced bands for pixel mode
+        in_pixels: Pixel vs language (inv freq) mode.
+        dtype: Output dtype.
+        device: Output device.
+
+    Returns:
+
+    """
+    sin_emb, cos_emb = build_fourier_pos_embed(
+        feat_shape,
+        bands=bands,
+        num_bands=dim // 4,
+        max_res=max_res,
+        temperature=temperature,
+        linear_bands=linear_bands,
+        in_pixels=in_pixels,
+        ref_feat_shape=ref_feat_shape,
+        device=device,
+        dtype=dtype,
+    )
+    num_spatial_dim = 1
+    # this would be much nicer as a .numel() call to torch.Size(), but torchscript sucks
+    for x in feat_shape:
+        num_spatial_dim *= x
+    sin_emb = sin_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1)
+    cos_emb = cos_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1)
+    return sin_emb, cos_emb
+
+def freq_bands(
+        num_bands: int,
+        temperature: float = 10000.,
+        step: int = 2,
+        device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands
+    bands = 1. / (temperature ** exp)
+    return bands
+
+    
+def pixel_freq_bands(
+        num_bands: int,
+        max_freq: float = 224.,
+        linear_bands: bool = True,
+        device: Optional[torch.device] = None,
+):
+    if linear_bands:
+        bands = torch.linspace(1.0, max_freq / 2, num_bands, dtype=torch.float32, device=device)
+    else:
+        bands = 2 ** torch.linspace(0, math.log(max_freq, 2) - 1, num_bands, dtype=torch.float32, device=device)
+    return bands * torch.pi
+
+def rot(x):
+    return torch.stack([-x[..., 1::2], x[..., ::2]], -1).reshape(x.shape)
+
+def apply_rot_embed_cat(x: torch.Tensor, emb):
+    sin_emb, cos_emb = emb.tensor_split(2, -1)
+    if sin_emb.ndim == 3:
+        return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze(1).expand_as(x)
+    return x * cos_emb + rot(x) * sin_emb
+
+# taken from https://github.com/huggingface/pytorch-image-models/blob/cb0e4391beedcc5ac3ae4bce16561b95c326f32c/timm/layers/pos_embed_sincos.py#L363
+class NomicVisionRotaryEmbeddingCat(nn.Module):
+    """ Rotary position embedding w/ concatenatd sin & cos
+
+    The following impl/resources were referenced for this impl:
+    * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py
+    * https://blog.eleuther.ai/rotary-embeddings/
+    """
+
+    def __init__(
+            self,
+            dim,
+            max_res=224,
+            temperature=10000,
+            in_pixels=True,
+            linear_bands: bool = False,
+            feat_shape: Optional[List[int]] = None,
+            ref_feat_shape: Optional[List[int]] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_res = max_res
+        self.temperature = temperature
+        self.in_pixels = in_pixels
+        self.feat_shape = feat_shape
+        self.ref_feat_shape = ref_feat_shape
+
+        if feat_shape is None:
+            # only cache bands
+            if in_pixels:
+                bands = pixel_freq_bands(
+                    dim // 4,
+                    float(max_res),
+                    linear_bands=linear_bands,
+                )
+            else:
+                bands = freq_bands(
+                    dim // 4,
+                    temperature=temperature,
+                    step=1,
+                )
+            self.register_buffer(
+                'bands',
+                bands,
+                persistent=False,
+            )
+            self.pos_embed = None
+        else:
+            # cache full sin/cos embeddings if shape provided up front
+            embeds = build_rotary_pos_embed(
+                feat_shape=feat_shape,
+                dim=dim,
+                max_res=max_res,
+                linear_bands=linear_bands,
+                in_pixels=in_pixels,
+                ref_feat_shape=self.ref_feat_shape,
+            )
+            self.bands = None
+            self.register_buffer(
+                'pos_embed',
+                torch.cat(embeds, -1),
+                persistent=False,
+            )
+
+    def get_embed(self, shape: Optional[List[int]] = None):
+        if self.bands is not None and shape is not None:
+            # rebuild embeddings every call, use if target shape changes
+            embeds = build_rotary_pos_embed(
+                shape,
+                self.bands,
+                in_pixels=self.in_pixels,
+                ref_feat_shape=self.ref_feat_shape,
+            )
+            return torch.cat(embeds, -1)
+        elif self.pos_embed is not None:
+            return self.pos_embed
+        else:
+            assert False, "get_embed() requires pre-computed pos_embed or valid shape w/ pre-computed bands"
+
+    def forward(self, x):
+        # assuming channel-first tensor where spatial dim are >= 2
+        pos_embed = self.get_embed(x.shape[2:])
+        return apply_rot_embed_cat(x, pos_embed)
+
+class NomicVisionPatchEmbeddings(nn.Module):
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__()
+        img_size = _pair(config.img_size)
+        patch_size = _pair(config.patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+        self.proj = nn.Linear(
+            config.num_channels * patch_size[0] * patch_size[1], config.n_embd, bias=config.patch_embed_bias
+        )
+
+        self.learned_pos_embedding = False
+        self.sinusoidal_pos_embedding = False
+        self.no_embed_class = getattr(config, "no_embed_class", False)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.n_embd)) if not getattr(config, "no_cls_token", False) else None
+        if config.learned_pos_embedding:
+            # this is the default in DINO
+            self.learned_pos_embedding = True
+            # hack for timm dinov2 with registers
+            num_patches = self.num_patches if getattr(config, "register_tokens", 0) > 0 else self.num_patches + 1
+            self.pos_embed = nn.Parameter(torch.randn(1, num_patches, config.n_embd) * 0.02) if getattr(config, "use_pos_embed", True) else None
+        elif getattr(config, "sinusoidal_pos_embedding", False):
+            self.sinusoidal_pos_embedding = True
+            if getattr(config, "use_pos_embed", True):
+                self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.n_embd), requires_grad=False)
+                pos_embed = get_2d_sincos_pos_embed(config.n_embd, self.grid_size[0], add_cls_token=True)
+                self.pos_embed.data.copy_(torch.from_numpy(pos_embed).to(self.pos_embed))
+            else:
+                self.pos_embed = None
+        else:
+            self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches + 1, config.n_embd) * 0.02) if getattr(config, "use_pos_embed", True) else None
+
+        if getattr(config, "register_tokens", 0) > 0:
+            self.reg_token = nn.Parameter(torch.randn(1, config.register_tokens, config.n_embd) * 0.02)
+        else:
+            self.reg_token = None 
+
+        if config.mask_token:
+            self.mask_token = nn.Parameter(torch.zeros(1, config.n_embd))
+
+        self.patch_dropout = nn.Identity()
+
+        if getattr(config, "use_rotary_pos_emb", False):
+            ref_feat_shape = getattr(config, "ref_feat_shape", None)
+            ref_feat_shape = to_2tuple(ref_feat_shape) if ref_feat_shape is not None else None
+            self.rope = NomicVisionRotaryEmbeddingCat(
+                config.n_embd // config.n_head,
+                in_pixels=False,
+                feat_shape=self.grid_size,
+                ref_feat_shape=ref_feat_shape,
+            )
+        else:
+            self.rope = None
+
+
+    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+        num_patches = embeddings.shape[1] - 1
+        num_positions = self.pos_embed.shape[1] - 1
+        if num_patches == num_positions and height == width:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = embeddings.shape[-1]
+        height = height // self.patch_size[0]
+        width = width // self.patch_size[1]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        height, width = height + 0.1, width + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            scale_factor=(height / math.sqrt(num_positions), width / math.sqrt(num_positions)),
+            mode="bicubic",
+            align_corners=False,
+        )
+        if int(height) != patch_pos_embed.shape[-2] or int(width) != patch_pos_embed.shape[-1]:
+            raise ValueError("Width or height does not match with the interpolated position embeddings")
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def forward(self, x):
+        # deepspeed case where the input is in fp32
+        if x.dtype != self.proj.weight.dtype:
+            x = x.to(dtype=self.proj.weight.dtype)
+
+        _, _, height, width = x.shape
+        x = self.proj(
+            rearrange(
+                x,
+                "b c (h p1) (w p2) -> b h w (c p1 p2)",
+                p1=self.patch_size[0],
+                p2=self.patch_size[1],
+            )
+        )
+        embeddings = rearrange(x, "b h w c -> b (h w) c")
+
+        to_cat = []
+        if self.cls_token is not None:
+            if self.sinusoidal_pos_embedding:
+                cls_token = self.cls_token + self.pos_embed[:, 0]
+                cls_token = cls_token.expand(embeddings.shape[0], -1, -1)
+                to_cat += [cls_token]
+            else:
+                cls_token = self.cls_token.expand(embeddings.shape[0], 1, -1)
+                to_cat += [cls_token]
+
+        if self.reg_token is not None:
+            to_cat += [self.reg_token.expand(embeddings.shape[0], -1, -1)]
+
+        rot_pos_embed = self.rope.get_embed() if self.rope is not None else None
+
+        if self.no_embed_class:
+            if self.learned_pos_embedding:
+                embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+            else:
+                if self.pos_embed is not None:
+                    embeddings = embeddings + self.pos_embed 
+            if to_cat:
+                embeddings = torch.cat(to_cat + [embeddings], dim=1)
+        else:
+            if to_cat:
+                embeddings = torch.cat(to_cat + [embeddings], dim=1)
+            if self.learned_pos_embedding:
+                if self.pos_embed is not None:
+                    embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+            else:
+                if self.pos_embed is not None:
+                    embeddings = embeddings + self.pos_embed 
+
+        embeddings = self.patch_dropout(embeddings)
+
+        return embeddings, rot_pos_embed
+
+
+class NomicBertEmbeddings(nn.Module):
+    def __init__(self, config):
+        """
+        If max_position_embeddings <= 0, there's no position embeddings
+        If type_vocab_size <= 0, there's no token type embeddings
+        """
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.max_position_embeddings = config.max_position_embeddings if config.rotary_emb_fraction <= 0 else 0
+        self.type_vocab_size = config.type_vocab_size
+        if self.max_position_embeddings > 0 and config.rotary_emb_fraction <= 0:
+            self.position_embeddings = nn.Embedding(
+                config.max_position_embeddings,
+                config.hidden_size,
+            )
+        if self.type_vocab_size > 0:
+            self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+    def forward(self, input_ids, position_ids=None, token_type_ids=None):
+        """
+        input_ids: (batch, seqlen)
+        position_ids: (batch, seqlen)
+        token_type_ids: (batch, seqlen)
+        """
+        batch_size, seqlen = input_ids.shape
+        embeddings = self.word_embeddings(input_ids)
+
+        if self.type_vocab_size > 0:
+            if token_type_ids is None:
+                token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=input_ids.device)
+            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+            embeddings = embeddings + token_type_embeddings
+
+        if self.max_position_embeddings > 0:
+            if position_ids is None:
+                position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
+            position_embeddings = self.position_embeddings(position_ids)
+            embeddings = embeddings + position_embeddings
+        return embeddings
+
+
+class NomicBertMLP(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        activation=F.gelu,
+        bias1=True,
+        bias2=True,
+        return_residual=False,
+        fused_bias_fc=False,
+    ):
+        super().__init__()
+        out_features = out_features if out_features is not None else in_features
+        hidden_features = hidden_features if hidden_features is not None else in_features * 4
+        self.return_residual = return_residual
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1)
+        approximate = "tanh" if activation in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] else "none"
+        self.activation = nn.GELU(approximate=approximate) if activation == "gelu" else activation
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)
+
+    def forward(self, x):
+        y = self.fc1(x)
+        y = self.activation(y)
+        y = self.fc2(y)
+        return y if not self.return_residual else (y, x)
+
+
+class NomciBertGatedMLP(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        activation=F.sigmoid,
+        bias1=True,
+        bias2=True,
+        multiple_of=256,
+        return_residual=False,
+        fused_bias_fc=True,
+        device=None,
+        dtype=None,
+        norm_layer=False,
+    ):
+        super().__init__()
+        out_features = out_features if out_features is not None else in_features
+        hidden_features = hidden_features if hidden_features is not None else int(8 * in_features / 3)
+        hidden_features = int((hidden_features + multiple_of - 1) // multiple_of * multiple_of)
+        self.return_residual = return_residual
+
+        self.fc11 = nn.Linear(in_features, hidden_features, bias=bias1)
+        self.fc12 = nn.Linear(in_features, hidden_features, bias=bias1)
+        self.activation = activation
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)
+        self.norm = nn.LayerNorm(hidden_features) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        y = self.fc11(x)
+        gate = self.fc12(x)
+        if self.activation == F.sigmoid:  # Special case for GLU
+            y = F.glu(torch.cat([y, gate], dim=-1), dim=-1)
+        else:
+            y = y * self.activation(gate)
+
+        # eva uses layer norm after the activation
+        y = self.norm(y)
+
+        y = self.fc2(y)
+        return y if not self.return_residual else (y, x)
+
+
+def rotate_half(x, interleaved=False):
+    if not interleaved:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    else:
+        x1, x2 = x[..., ::2], x[..., 1::2]
+        return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
+
+
+def apply_rotary_emb(x, cos, sin, offset=0, interleaved=False):
+    """
+    x: (batch_size, seqlen, nheads, headdim)
+    cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
+    """
+    ro_dim = cos.shape[-1] * 2
+    assert ro_dim <= x.shape[-1]
+    cos, sin = (
+        cos[offset : offset + x.shape[1]],
+        sin[offset : offset + x.shape[1]],
+    )
+    cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
+    sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
+    return torch.cat(
+        [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]],
+        dim=-1,
+    )
+
+
+class NomicBertRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        base=10000.0,
+        interleaved=False,
+        scale_base=None,
+        pos_idx_in_fp32=True,
+        device=None,
+    ):
+        """
+        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+            of 1st half and 2nd half (GPT-NeoX style).
+        pos_idx_in_fp32: if True, the position indices [0.0, ..., seqlen - 1] are in fp32,
+            otherwise they might be in lower precision.
+            This option was added because previously (before 2023-07-02), when we construct
+            the position indices, we use the dtype of self.inv_freq. In most cases this would
+            be fp32, but if the model is trained in pure bf16 (not mixed precision), then
+            self.inv_freq would be bf16, and the position indices are also in bf16.
+            Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the
+            embeddings for some positions will coincide.
+            To maintain compatibility with models previously trained in pure bf16,
+            we add this option.
+        """
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+        # Generate and save the inverse frequency buffer (non trainable)
+        inv_freq = self._compute_inv_freq(device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.interleaved = interleaved
+        self.scale_base = scale_base
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
+
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cos_k_cached = None
+        self._sin_k_cached = None
+
+    def _compute_inv_freq(self, device=None):
+        return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        # Reset the tables if the sequence length has changed,
+        # if we're on a new device (possibly due to tracing for instance),
+        # or if we're switching from inference mode to training
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
+            # And the output of arange can be quite large, so bf16 would lose a lot of precision.
+            # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                # We want fp32 here as well since inv_freq will be multiplied with t, and the output
+                # will be large. Having it in bf16 will lose a lot of precision and cause the
+                # cos & sin output to change significantly.
+                # We want to recompute self.inv_freq if it was not loaded in fp32
+                if self.inv_freq.dtype != torch.float32:
+                    inv_freq = self._compute_inv_freq(device=device)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                inv_freq = self.inv_freq
+            # Don't do einsum, it converts fp32 to fp16 under AMP
+            # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            freqs = torch.outer(t, inv_freq)
+            self._cos_cached = torch.cos(freqs).to(dtype)
+            self._sin_cached = torch.sin(freqs).to(dtype)
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        kv: Optional[torch.Tensor] = None,
+        seqlen_offset: Union[int, torch.Tensor] = 0,
+        max_seqlen: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        qkv: (batch, seqlen, 3, nheads, headdim) if kv is none,
+             else it's just q of shape (batch, seqlen, nheads, headdim)
+        kv: (batch, seqlen, 2, nheads, headdim)
+        seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount.
+            Most commonly used in inference when we have KV cache.
+            If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one
+            should pass in max_seqlen, which will update the cos / sin cache up to that length.
+        Apply rotary embedding *inplace* to qkv and / or kv.
+        """
+        seqlen = qkv.shape[1]
+        if seqlen > self._seq_len_cached:
+            self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
+        elif max_seqlen is not None:
+            self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
+        elif isinstance(seqlen_offset, int):
+            self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
+
+        q_rot = apply_rotary_emb(qkv[:, :, 0], self._cos_cached, self._sin_cached, seqlen_offset, self.interleaved)
+        k_rot = apply_rotary_emb(qkv[:, :, 1], self._cos_cached, self._sin_cached, seqlen_offset, self.interleaved)
+        return torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
+
+
+class NomicBertDynamicNTKRotaryEmbedding(NomicBertRotaryEmbedding):
+    def __init__(self, rotary_scaling_factor, max_position_embeddings, **kwargs):
+        super().__init__(**kwargs)
+        self.rotary_scaling_factor = rotary_scaling_factor
+        self.max_position_embeddings = max_position_embeddings
+
+    def _compute_inv_freq(self, base=None, device=None):
+        if base is None:
+            base = self.base
+        return 1.0 / (base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        # Reset the tables if the sequence length has changed,
+        # if we're on a new device (possibly due to tracing for instance),
+        # or if we're switching from inference mode to training
+        if seqlen > self.max_position_embeddings:
+            base = self.base * (
+                (self.rotary_scaling_factor * seqlen / self.max_position_embeddings) - (self.rotary_scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = self._compute_inv_freq(base=base, device=device)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
+            # And the output of arange can be quite large, so bf16 would lose a lot of precision.
+            # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                # We want fp32 here as well since inv_freq will be multiplied with t, and the output
+                # will be large. Having it in bf16 will lose a lot of precision and cause the
+                # cos & sin output to change significantly.
+                # We want to recompute self.inv_freq if it was not loaded in fp32
+                if self.inv_freq.dtype != torch.float32:
+                    if seqlen > self.max_position_embeddings:
+                        base = self.base * (
+                            (self.scaling_factor * seqlen / self.max_position_embeddings) - (self.scaling_factor - 1)
+                        ) ** (self.dim / (self.dim - 2))
+                    else:
+                        base = self.base
+                    inv_freq = self._compute_inv_freq(device=device, base=base)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                inv_freq = self.inv_freq
+            # Don't do einsum, it converts fp32 to fp16 under AMP
+            # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            freqs = torch.outer(t, inv_freq)
+            if self.scale is None:
+                self._cos_cached = torch.cos(freqs).to(dtype)
+                self._sin_cached = torch.sin(freqs).to(dtype)
+            else:
+                power = (
+                    torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
+                ) / self.scale_base
+                scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+                # We want the multiplication by scale to happen in fp32
+                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+
+class NomicBertAttention(nn.Module):
+    """Multi-head self-attention and cross-attention"""
+
+    def __init__(
+        self,
+        config,
+    ) -> None:
+        """
+        num_heads_kv: can be used to toggle MQA / GQA. If None, use num_heads.
+        return_residual: whether to return the input x along with the output. This is for
+            performance reason: for post-norm architecture, returning the input allows us
+            to fuse the backward of nn.Linear with the residual connection.
+        """
+        super().__init__()
+        self.embed_dim = config.n_embd
+        self.use_flash_attn = config.use_flash_attn
+        self.fused_bias_fc = config.fused_bias_fc
+
+        self.num_heads = config.n_head
+        self.num_heads_kv = config.num_heads_kv if getattr(config, "num_heads_kv", None) is not None else self.num_heads
+        assert self.embed_dim % self.num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.head_dim = self.embed_dim // self.num_heads
+        # we don't really support mqa / gqa for now
+        qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_heads_kv)
+
+        self.register_buffer(
+            "norm_factor",
+            torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
+            persistent=False,
+        )
+
+        self.rotary_emb_dim = self.head_dim * config.rotary_emb_fraction
+        if self.rotary_emb_dim > 0:
+            if getattr(config, "rotary_scaling_factor", None):
+                self.rotary_emb = NomicBertDynamicNTKRotaryEmbedding(
+                    dim=self.rotary_emb_dim,
+                    base=config.rotary_emb_base,
+                    scale_base=config.rotary_emb_scale_base,
+                    interleaved=config.rotary_emb_interleaved,
+                    rotary_scaling_factor=config.rotary_scaling_factor,
+                    max_position_embeddings=config.max_trained_positions,
+                )
+            else:
+                self.rotary_emb = NomicBertRotaryEmbedding(
+                    dim=self.rotary_emb_dim,
+                    base=config.rotary_emb_base,
+                    scale_base=config.rotary_emb_scale_base,
+                    interleaved=config.rotary_emb_interleaved,
+                )
+            # bug in xformers: https://github.com/facebookresearch/xformers/issues/841
+            # uses the head dimension instead of the sequence dimension
+            self.rotary_head_dim = getattr(config, "rotary_head_dim", False)
+
+        self.Wqkv = nn.Linear(self.embed_dim, qkv_dim, bias=config.qkv_proj_bias)
+
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
+        self.causal = config.causal
+        self.drop = nn.Dropout(config.attn_pdrop)
+        self.num_prefix_tokens = max(getattr(config, "register_tokens", 1), 1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        is_padded_inputs: Optional[bool] = True,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seq_len: Optional[int] = None,
+        rope: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+
+        has_layer_past = past_key_value is not None
+
+        if has_layer_past:
+            past_key_value = past_key_value[0]
+            past_len = past_key_value[1]
+        else:
+            past_len = 0
+
+        qkv = self.Wqkv(hidden_states)
+        qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
+
+        past_key_value = (past_key_value, past_len + qkv.size(1)) if use_cache else None
+
+        if self.rotary_emb_dim > 0:
+            if self.rotary_head_dim:
+                qkv = rearrange(qkv, "b s three h d -> b h three s d")
+            qkv = self.rotary_emb(qkv, seqlen_offset=past_len)
+
+            if self.rotary_head_dim:
+                qkv = rearrange(qkv, "b h three s d -> b s three h d")
+        elif rope is not None:
+            q, k, v = qkv.permute(0, 3, 1, 2, 4).unbind(dim=-2)
+            q = torch.cat([q[:, :, :self.num_prefix_tokens], apply_rot_embed_cat(q[:, :, self.num_prefix_tokens:], rope)], dim=2).type_as(q)
+            k = torch.cat([k[:, :, :self.num_prefix_tokens], apply_rot_embed_cat(k[:, :, self.num_prefix_tokens:], rope)], dim=2).type_as(q)
+
+            qkv = torch.stack([q, k, v], dim=-2)
+            qkv = rearrange(qkv, "b h s three d -> b s three h d")
+
+        query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+
+        attention_scores = torch.matmul(query, key.transpose(-1, -2)) / self.norm_factor
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+
+        attentions_probs = F.softmax(attention_scores, dim=-1)
+        attentions_probs = self.drop(attentions_probs)
+
+        attn_output = torch.matmul(attentions_probs, value)
+        attn_output = rearrange(attn_output.permute(0, 2, 1, 3), "... h d -> ... (h d)")
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output
+
+
+class NomicBertBlock(NomicBertPreTrainedModel):
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(config=config)
+        self.prenorm = config.prenorm
+        self.fused_dropout_add_ln = config.fused_dropout_add_ln
+
+        self.attn = NomicBertAttention(config)
+        activation = (
+            F.sigmoid
+            if config.activation_function == "glu"
+            else (F.silu if config.activation_function == "swiglu" else F.gelu)
+        )
+        if config.activation_function in ["glu", "swiglu", "geglu"]:
+            self.mlp = NomciBertGatedMLP(
+                config.n_embd,
+                hidden_features=config.n_inner,
+                bias1=config.mlp_fc1_bias,
+                bias2=config.mlp_fc2_bias,
+                activation=activation,
+                fused_bias_fc=config.fused_bias_fc,
+                norm_layer=getattr(config, "norm_mlp", False),
+            )
+        else:
+            self.mlp = NomicBertMLP(
+                config.n_embd,
+                hidden_features=config.n_inner,
+                bias1=config.mlp_fc1_bias,
+                bias2=config.mlp_fc2_bias,
+                activation=activation,
+                fused_bias_fc=config.fused_bias_fc,
+            )
+
+        self.dropout1 = nn.Dropout(config.resid_pdrop)
+        self.norm1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.norm2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.dropout2 = nn.Dropout(config.resid_pdrop)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states2: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        is_padded_inputs: Optional[bool] = True,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seq_len: Optional[int] = None,
+        rope: Optional[torch.Tensor] = None,
+    ):
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: if postnorm, residual=None, If prenorm, hidden_states = Attn/MLP(LN(residual))
+            mixer_subset: for cross-attention only. If not None, will take a subset of x
+                before applying the query projection. Useful for e.g., ViT where we only care
+                about the CLS token in the last layer.
+        """
+        if self.prenorm:
+            dropped = self.dropout1(hidden_states)
+            residual = (dropped + residual) if residual is not None else dropped
+            hidden_states = self.norm1(residual.to(dtype=self.norm1.weight.dtype))
+            hidden_states = self.attn(
+                hidden_states,
+                attention_mask=attention_mask,
+                is_padded_inputs=is_padded_inputs,
+                cu_seqlens=cu_seqlens,
+                max_seq_len=max_seq_len,
+                rope=rope,
+            )
+
+            dropped = self.dropout2(hidden_states)
+            residual = (dropped + residual) if residual is not None else dropped
+            hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
+            hidden_states = self.mlp(hidden_states)
+
+            return hidden_states, None, residual
+        else:
+            assert residual is None
+            attn_outputs = self.attn(
+                hidden_states,
+                attention_mask=attention_mask,
+                is_padded_inputs=is_padded_inputs,
+                cu_seqlens=cu_seqlens,
+                max_seq_len=max_seq_len,
+                rope=rope,
+            )
+            hidden_states = self.norm1((self.dropout1(attn_outputs) + hidden_states).to(dtype=self.norm1.weight.dtype))
+            mlp_out = self.mlp(hidden_states)
+
+            hidden_states = self.norm2((self.dropout2(mlp_out) + hidden_states).to(dtype=self.norm2.weight.dtype))
+            return hidden_states, None, None
+
+
+class NomicBertEncoder(nn.Module):
+    def __init__(self, config: GPT2Config):
+        super().__init__()
+        self.layers = nn.ModuleList([NomicBertBlock(config) for _ in range(config.n_layer)])
+        self.gradient_checkpointing = False
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        is_padded_inputs: Optional[bool] = True,
+        rope: Optional[torch.Tensor] = None,
+    ):
+        """If subset_mask is not None, we only want output for the subset of the sequence.
+        This means that we only compute the last layer output for these tokens.
+        subset_mask: (batch, seqlen), dtype=torch.bool
+        """
+        hidden_states2 = None
+        residual = None
+
+        for _, layer in enumerate(self.layers):
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states, hidden_states2, residual = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer),
+                    hidden_states,
+                    hidden_states2,
+                    residual,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    is_padded_inputs,
+                    output_attentions,
+                    use_cache,
+                    None,
+                    None,
+                    rope,
+                    # if you freeze ANY layers, you need `use_reentrant=False`
+                    # https://github.com/huggingface/transformers/issues/21381
+                    # https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/7
+                    use_reentrant=False,
+                )
+
+            else:
+                hidden_states, hidden_states2, residual = layer(
+                    hidden_states,
+                    hidden_states2,
+                    residual,
+                    attention_mask,
+                    position_ids,
+                    None,
+                    is_padded_inputs,
+                    output_attentions,
+                    use_cache,
+                    rope=rope,
+                )
+        return hidden_states
+
+
+class NomicBertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.n_embd, config.n_embd)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states, pool=True):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0] if pool else hidden_states
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class NomicBertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.n_embd, config.n_embd, bias=config.mlp_fc1_bias)
+        approximate = "tanh" if config.activation_function in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] else "none"
+        if config.activation_function == "swiglu":
+            self.transform_act_fn = F.silu
+        else:
+            self.transform_act_fn = nn.GELU(approximate=approximate)
+
+        self.layer_norm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+
+        return hidden_states
+
+
+class NomicBertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.transform = NomicBertPredictionHeadTransform(config)
+
+        self.decoder = nn.Linear(config.n_embd, config.vocab_size, bias=config.mlp_fc1_bias)
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+class NomicBertPreTrainingHeads(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = NomicBertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class NomicBertModel(NomicBertPreTrainedModel):
+    def __init__(self, config: GPT2Config, add_pooling_layer=True):
+        super().__init__(config)
+        self.pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
+        if config.vocab_size % self.pad_vocab_size_multiple != 0:
+            config.vocab_size += self.pad_vocab_size_multiple - (config.vocab_size % self.pad_vocab_size_multiple)
+
+        assert config.activation_function in [
+            "gelu",
+            "gelu_new",
+            "gelu_fast",
+            "gelu_pytorch_tanh",
+            "swiglu",
+            "geglu",
+            "glu",
+        ]
+
+        self.embeddings = NomicBertEmbeddings(config)
+        self.emb_drop = nn.Dropout(config.resid_pdrop)
+        self.emb_ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.encoder = NomicBertEncoder(config)
+        self.pooler = NomicBertPooler(config) if add_pooling_layer else None
+
+        self.apply(partial(_init_weights, initializer_range=config.initializer_range))
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        position_ids=None,
+        token_type_ids=None,
+        return_dict=None,
+        matryoshka_dim=None,
+    ):
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+        hidden_states = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)
+        hidden_states = self.emb_ln(hidden_states)
+        hidden_states = self.emb_drop(hidden_states)
+
+        attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape)
+        sequence_output = self.encoder(hidden_states, attention_mask=attention_mask, return_dict=return_dict)
+
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        if matryoshka_dim:
+            sequence_output = sequence_output[:, :matryoshka_dim]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+        )
+
+
+class NomicBertForPreTraining(NomicBertPreTrainedModel):
+    _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
+
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+
+        self.bert = NomicBertModel(config, add_pooling_layer=getattr(config, "add_pooling_layer", False))
+        self.cls = NomicBertPreTrainingHeads(config)
+        self.mlm_loss = nn.CrossEntropyLoss()
+
+        # Initialize weights and apply final processing
+        self.apply(partial(_init_weights, initializer_range=config.initializer_range))
+        self.tie_weights()
+
+    def tie_weights(self):
+        self.cls.predictions.decoder.weight = self.bert.embeddings.word_embeddings.weight
+
+    def forward(
+        self,
+        input_ids,
+        position_ids=None,
+        token_type_ids=None,
+        attention_mask=None,
+        labels=None,
+    ):
+        """
+        If labels are provided, they must be -100 for masked out tokens (as specified in the attention
+        mask).
+        Outputs:
+            if `labels` and `next_sentence_label` are not `None`:
+                Outputs the total_loss which is the sum of the masked language modeling loss and the next
+                sentence classification loss.
+            if `labels` or `next_sentence_label` is `None`:
+                Outputs a tuple comprising
+                - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
+                - the next sentence classification logits of shape [batch_size, 2].
+
+        """
+        outputs = self.bert(
+            input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask.bool() if attention_mask is not None else None,
+        )
+        sequence_output, _ = outputs.last_hidden_state, outputs.pooler_output
+
+        prediction_scores = self.cls(sequence_output)
+
+        total_loss = None
+        if labels is not None:
+            masked_lm_loss = self.mlm_loss(
+                rearrange(prediction_scores, "... v -> (...) v"),
+                rearrange(labels, "... -> (...)"),
+            )
+            total_loss = masked_lm_loss.float()
+
+        return MaskedLMOutput(
+            loss=total_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=None,
+        )
+
+
+class NomicBertForSequenceClassification(NomicBertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = NomicBertModel(config)
+        classifier_dropout = getattr(config, "classifier_dropout", config.embd_pdrop)
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.n_embd, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.bert(
+            input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask.bool() if attention_mask is not None else None,
+        )
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = nn.MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = nn.BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+def hf_vit_config_to_vit_config(vit_config: ViTConfig) -> GPT2Config:
+    return GPT2Config(
+        n_embd=vit_config.hidden_size,
+        n_layer=vit_config.num_hidden_layers,
+        n_head=vit_config.num_attention_heads,
+        n_inner=vit_config.intermediate_size,
+        activation_function=vit_config.hidden_act,
+        vocab_size=0, # no vocab since using patches
+        n_positions=0,  # No absolute position embedding
+        resid_pdrop=0.0,  # No dropout
+        embd_pdrop=getattr(vit_config, "dropout", 0.0),
+        attn_pdrop=vit_config.attention_probs_dropout_prob,
+        layer_norm_epsilon=vit_config.layer_norm_eps,
+        initializer_range=vit_config.initializer_range,
+        bos_token_id=None,
+        eos_token_id=None,
+        # These are new arguments not in the original GPT2Config
+        drop_path_rate=0.0,
+        # Why is there double layer norm??
+        prepre_layernom=False,
+        layer_scale=False,
+        layer_scale_init=None,
+        img_size=vit_config.image_size,
+        patch_size=vit_config.patch_size,
+        num_channels=vit_config.num_channels,
+        prenorm=True,
+        parallel_block=False,
+        parallel_block_tied_norm=False,
+        rotary_emb_fraction=0,
+        tie_word_embeddings=False,
+        fused_dropout_add_ln=True,
+        fused_bias_fc=True,
+        patch_embed_bias=True,
+        use_flash_attn=True,
+        qkv_proj_bias=True,
+        mlp_fc1_bias=getattr(vit_config, "mlp_fc1_bias", True),
+        mlp_fc2_bias=getattr(vit_config, "mlp_fc2_bias", True),
+        use_rms_norm=False,
+        causal=False,
+        hidden_features_scaling_factor=1.0,
+        mask_token=False,
+        learned_pos_embedding=False,
+        patch_dropout=0,
+        sinusoidal_pos_embedding=vit_config.model_type == "vit_mae"
+    )
+
+    
+class NomicAttentionPooling(nn.Module):
+    def __init__(
+        self,
+        config
+    ):
+        super().__init__()
+        self.embed_dim = config.n_embd
+        self.use_flash_attn = config.use_flash_attn
+        self.fused_bias_fc = config.fused_bias_fc
+
+        self.num_heads = config.n_head
+        self.num_heads_kv = config.num_heads_kv if getattr(config, "num_heads_kv", None) is not None else self.num_heads
+        assert self.embed_dim % self.num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.head_dim = self.embed_dim // self.num_heads
+        # we don't really support mqa / gqa for now
+        kv_dim = 2 * self.head_dim * self.num_heads_kv
+
+        self.register_buffer(
+            "norm_factor",
+            torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
+            persistent=False,
+        )
+
+        self.Wq = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
+        self.Wkv = nn.Linear(self.embed_dim, kv_dim, bias=config.qkv_proj_bias)
+        
+        self.latent = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
+        self.causal = config.causal
+        self.drop = nn.Dropout(config.attn_pdrop)
+
+    def init_weights(self):
+        trunc_normal_tf_(self.latent, std=self.embed_dim ** -0.5)
+
+    def forward(
+        self,
+        kv,
+        attention_mask=None,
+        cu_seqlens_k=None,
+        max_seqlen_k=None,
+        is_padded_inputs: Optional[bool] = True,
+        output_attentions: bool = False,
+    ):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            q: The tensor containing the query. (B, Sq, H, D)
+            kv: The tensor containing the key and value. (B, Sk, 2, H_k, D)
+            causal: if passed, will override self.causal
+            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+                of the sequences in the batch, used to index into q.
+            max_seqlen: int. Maximum sequence length in the batch of q.
+            cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+                of the sequences in the batch, used to index into kv.
+            max_seqlen_k: int. Maximum sequence length in the batch of k and v.
+        """
+        q_latent = self.latent.expand(kv.size(0), -1, -1)
+        q = self.Wq(q_latent)
+        bsz, q_len, h_size = q.shape
+        kv = self.Wkv(kv)
+        query = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
+
+        key, value = kv[:, :, 0], kv[:, :, 1]
+
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+
+        attention_scores = torch.matmul(query, key.transpose(-1, -2)) / self.norm_factor
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+
+        attentions_probs = F.softmax(attention_scores, dim=-1)
+        attentions_probs = self.drop(attentions_probs)
+
+        attn_output = torch.matmul(attentions_probs, value)
+        attn_output = rearrange(attn_output.permute(0, 2, 1, 3), "... h d -> ... (h d)")
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output
+
+        
+class NomicMultiHeadAttentionPooling(nn.Module):
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__()
+        self.prenorm = config.prenorm
+        self.fused_dropout_add_ln = config.fused_dropout_add_ln
+
+        self.attn = NomicAttentionPooling(config)
+        activation = (
+            F.sigmoid
+            if config.activation_function == "glu"
+            else (F.silu if config.activation_function == "swiglu" else F.gelu)
+        )
+        if config.activation_function in ["glu", "swiglu", "geglu"]:
+            self.mlp = NomciBertGatedMLP(
+                config.n_embd,
+                hidden_features=config.n_inner,
+                bias1=config.mlp_fc1_bias,
+                bias2=config.mlp_fc2_bias,
+                activation=activation,
+                fused_bias_fc=config.fused_bias_fc,
+            )
+        else:
+            self.mlp = NomicBertMLP(
+                config.n_embd,
+                hidden_features=config.n_inner,
+                bias1=config.mlp_fc1_bias,
+                bias2=config.mlp_fc2_bias,
+                activation=activation,
+                fused_bias_fc=config.fused_bias_fc,
+            )
+
+        self.dropout1 = nn.Dropout(config.resid_pdrop)
+        self.norm1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.dropout2 = nn.Dropout(config.resid_pdrop)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: if postnorm, residual=None, If prenorm, hidden_states = Attn/MLP(LN(residual))
+            mixer_subset: for cross-attention only. If not None, will take a subset of x
+                before applying the query projection. Useful for e.g., ViT where we only care
+                about the CLS token in the last layer.
+        """
+
+        attn_outputs = self.attn(
+                hidden_states,
+                attention_mask=attention_mask,
+            )
+
+        normed = self.norm1(attn_outputs)
+        hidden_states = hidden_states + self.mlp(normed)
+
+        return hidden_states
+
+class NomicVisionPreTrainedModel(PreTrainedModel):
+    """An abstract class to handle weights initialization and
+    a simple interface for dowloading and loading pretrained models.
+    """
+
+    config_class = NomicBertConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config)
+        if not isinstance(config, GPT2Config):
+            raise ValueError(
+                "Parameter config in `{}(config)` should be an instance of class `GPT2Config`. "
+                "To create a model from a Google pretrained model use "
+                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    self.__class__.__name__, self.__class__.__name__
+                )
+            )
+        self.config = config
+
+class NomicVisionModel(NomicVisionPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = NomicVisionPatchEmbeddings(config)
+        self.layers = nn.ModuleList([NomicBertBlock(config) for _ in range(config.n_layer)])
+
+        self.selector = NomicMultiHeadAttentionPooling(config)
+
+        self.global_pool = getattr(config, "global_pool", None)
+        self.num_prefix_tokens = (1 if not getattr(config, "no_cls_token", False) else 0) + getattr(config, "register_tokens", 0)
+
+        self.apply(partial(_init_weights, initializer_range=config.initializer_range))
+
+    def forward(
+        self,
+        pixel_values,
+        attention_mask=None,
+        position_ids=None,
+        token_type_ids=None,
+        return_dict=None,
+        matryoshka_dim=None,
+    ):
+        embeddings, rope = self.embeddings(pixel_values)
+
+        original_dtype = embeddings.dtype
+
+        hidden_states = embeddings 
+        # unused but easier to pass to gradient checkpointing as words
+        residual = None
+        for layer in self.layers:
+            # need to pass none for backwards compatability
+            hidden_states, _, residual = layer(hidden_states, None, residual=residual, is_padded_inputs=False, rope=rope)
+
+        hidden_states = hidden_states + residual
+        if self.global_pool == "avg":
+            hidden_states = hidden_states[:, self.num_prefix_tokens:].mean(dim=1)
+
+        pooled_output = self.selector(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=pooled_output,
+            hidden_states=hidden_states,
+        )

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 8192,
+  "do_lower_case": false
+}

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,62 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "max_length": 8192,
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}