Upload folder using huggingface_hub
Browse files- .gitattributes +1 -0
- 0_CollinsSTWrapper/config.json +18 -0
- 0_CollinsSTWrapper/model.safetensors +3 -0
- 0_CollinsSTWrapper/special_tokens_map.json +7 -0
- 0_CollinsSTWrapper/tokenizer.json +0 -0
- 0_CollinsSTWrapper/tokenizer_config.json +56 -0
- 0_CollinsSTWrapper/vocab.txt +0 -0
- README.md +134 -3
- collins_sts_comparison.pdf +0 -0
- collins_sts_comparison.png +3 -0
- config_sentence_transformers.json +14 -0
- modeling_hf.py +212 -0
- modules.json +8 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
collins_sts_comparison.png filter=lfs diff=lfs merge=lfs -text
|
0_CollinsSTWrapper/config.json
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"CollinsModel"
|
| 4 |
+
],
|
| 5 |
+
"attention_probs_dropout_prob": 0.1,
|
| 6 |
+
"dtype": "float32",
|
| 7 |
+
"hash_seed": 42,
|
| 8 |
+
"hidden_dropout_prob": 0.1,
|
| 9 |
+
"hidden_size": 256,
|
| 10 |
+
"intermediate_size": 1024,
|
| 11 |
+
"max_position_embeddings": 512,
|
| 12 |
+
"model_type": "collins",
|
| 13 |
+
"num_attention_heads": 8,
|
| 14 |
+
"num_buckets": 2048,
|
| 15 |
+
"num_hidden_layers": 3,
|
| 16 |
+
"transformers_version": "4.57.1",
|
| 17 |
+
"vocab_size": 30522
|
| 18 |
+
}
|
0_CollinsSTWrapper/model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c703775717ef9e52d8dba2441f5df63020ba0ead4e22e070e055f54890624a3d
|
| 3 |
+
size 12497664
|
0_CollinsSTWrapper/special_tokens_map.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cls_token": "[CLS]",
|
| 3 |
+
"mask_token": "[MASK]",
|
| 4 |
+
"pad_token": "[PAD]",
|
| 5 |
+
"sep_token": "[SEP]",
|
| 6 |
+
"unk_token": "[UNK]"
|
| 7 |
+
}
|
0_CollinsSTWrapper/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
0_CollinsSTWrapper/tokenizer_config.json
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"added_tokens_decoder": {
|
| 3 |
+
"0": {
|
| 4 |
+
"content": "[PAD]",
|
| 5 |
+
"lstrip": false,
|
| 6 |
+
"normalized": false,
|
| 7 |
+
"rstrip": false,
|
| 8 |
+
"single_word": false,
|
| 9 |
+
"special": true
|
| 10 |
+
},
|
| 11 |
+
"100": {
|
| 12 |
+
"content": "[UNK]",
|
| 13 |
+
"lstrip": false,
|
| 14 |
+
"normalized": false,
|
| 15 |
+
"rstrip": false,
|
| 16 |
+
"single_word": false,
|
| 17 |
+
"special": true
|
| 18 |
+
},
|
| 19 |
+
"101": {
|
| 20 |
+
"content": "[CLS]",
|
| 21 |
+
"lstrip": false,
|
| 22 |
+
"normalized": false,
|
| 23 |
+
"rstrip": false,
|
| 24 |
+
"single_word": false,
|
| 25 |
+
"special": true
|
| 26 |
+
},
|
| 27 |
+
"102": {
|
| 28 |
+
"content": "[SEP]",
|
| 29 |
+
"lstrip": false,
|
| 30 |
+
"normalized": false,
|
| 31 |
+
"rstrip": false,
|
| 32 |
+
"single_word": false,
|
| 33 |
+
"special": true
|
| 34 |
+
},
|
| 35 |
+
"103": {
|
| 36 |
+
"content": "[MASK]",
|
| 37 |
+
"lstrip": false,
|
| 38 |
+
"normalized": false,
|
| 39 |
+
"rstrip": false,
|
| 40 |
+
"single_word": false,
|
| 41 |
+
"special": true
|
| 42 |
+
}
|
| 43 |
+
},
|
| 44 |
+
"clean_up_tokenization_spaces": false,
|
| 45 |
+
"cls_token": "[CLS]",
|
| 46 |
+
"do_lower_case": true,
|
| 47 |
+
"extra_special_tokens": {},
|
| 48 |
+
"mask_token": "[MASK]",
|
| 49 |
+
"model_max_length": 512,
|
| 50 |
+
"pad_token": "[PAD]",
|
| 51 |
+
"sep_token": "[SEP]",
|
| 52 |
+
"strip_accents": null,
|
| 53 |
+
"tokenize_chinese_chars": true,
|
| 54 |
+
"tokenizer_class": "BertTokenizer",
|
| 55 |
+
"unk_token": "[UNK]"
|
| 56 |
+
}
|
0_CollinsSTWrapper/vocab.txt
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
README.md
CHANGED
|
@@ -1,3 +1,134 @@
|
|
| 1 |
-
---
|
| 2 |
-
|
| 3 |
-
--
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
tags:
|
| 3 |
+
- sentence-transformers
|
| 4 |
+
- sentence-similarity
|
| 5 |
+
- feature-extraction
|
| 6 |
+
- dense
|
| 7 |
+
- loss:MultipleNegativesRankingLoss
|
| 8 |
+
pipeline_tag: sentence-similarity
|
| 9 |
+
library_name: sentence-transformers
|
| 10 |
+
license: apache-2.0
|
| 11 |
+
language:
|
| 12 |
+
- en
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
# NoesisLab/Collins-Embedding-3M
|
| 16 |
+
|
| 17 |
+
A **3M-parameter** sentence embedding model built on **2-Universal Hash encoding + RoPE positional encoding**, trained on AllNLI triplets with MultipleNegativesRankingLoss.
|
| 18 |
+
|
| 19 |
+
The core insight: replace the vocabulary embedding table — the single largest cost in any transformer — with a 2-Universal Hash function that maps token IDs into a fixed-size bucket space in O(1) time. No lookup table. No gradient-heavy embedding matrix.
|
| 20 |
+
|
| 21 |
+
> Released 2026 by [NoesisLab](https://huggingface.co/NoesisLab).
|
| 22 |
+
|
| 23 |
+
---
|
| 24 |
+
|
| 25 |
+
## Quick Start
|
| 26 |
+
|
| 27 |
+
```python
|
| 28 |
+
from sentence_transformers import SentenceTransformer
|
| 29 |
+
|
| 30 |
+
model = SentenceTransformer("NoesisLab/Collins-Embedding-3M")
|
| 31 |
+
embeddings = model.encode(["Hello world", "Hi there"])
|
| 32 |
+
similarities = model.similarity(embeddings[0], embeddings[1])
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
---
|
| 36 |
+
|
| 37 |
+
## Architecture: Why Hashing Works
|
| 38 |
+
|
| 39 |
+
```
|
| 40 |
+
Token ID ──► h(x) = ((ax + b) mod p) mod B ──► bucket index
|
| 41 |
+
↑
|
| 42 |
+
Sign Hash: φ(x) = sign((cx + d) mod p)
|
| 43 |
+
resolves collision ambiguity during training
|
| 44 |
+
```
|
| 45 |
+
|
| 46 |
+
The **sign hash** acts as a per-token polarity signal. Under strong contrastive supervision, the model learns to disentangle hash collisions — tokens that share a bucket but carry different semantics get separated via their sign channel. The Chernoff Bound guarantees that the sign channel suppresses collision noise under sufficient supervision signal.
|
| 47 |
+
|
| 48 |
+
**Time complexity vs. standard embedding:**
|
| 49 |
+
|
| 50 |
+
| Operation | Standard Embedding | Collins Hash |
|
| 51 |
+
|---|---|---|
|
| 52 |
+
| Token → vector | O(1) table lookup | O(1) arithmetic |
|
| 53 |
+
| Memory (vocab) | O(V × d) | O(B × d), B ≪ V |
|
| 54 |
+
| Gradient flow | Dense, full vocab | Sparse, bucket-local |
|
| 55 |
+
| Cold-start | Requires pretraining | Random init viable |
|
| 56 |
+
|
| 57 |
+
With V = 30522 and B = 512, Collins uses ~60× fewer parameters for the token encoding stage alone.
|
| 58 |
+
|
| 59 |
+
**Cache efficiency**: At 3M total parameters, the entire model fits in GPU L2 cache during inference. Standard MiniLM models (15–22M) cannot achieve this, resulting in 1–2 orders of magnitude lower inference latency at equivalent semantic accuracy.
|
| 60 |
+
|
| 61 |
+
---
|
| 62 |
+
|
| 63 |
+
## MTEB Benchmark Results
|
| 64 |
+
|
| 65 |
+
| Task | cosine_spearman |
|
| 66 |
+
|---|---|
|
| 67 |
+
| STS12 | 0.6038 |
|
| 68 |
+
| STS13 | 0.5952 |
|
| 69 |
+
| STS14 | 0.6186 |
|
| 70 |
+
| **STSBenchmark** | **0.7114** |
|
| 71 |
+
|
| 72 |
+
---
|
| 73 |
+
|
| 74 |
+
## Full Baseline Comparison
|
| 75 |
+
|
| 76 |
+

|
| 77 |
+
|
| 78 |
+
*Left: STSBenchmark Spearman score. Right: score per million parameters (efficiency). White labels inside bars = parameter count.*
|
| 79 |
+
|
| 80 |
+
| Model | Type | Params | STSB Spearman | Score / M params |
|
| 81 |
+
|---|---|---|---|---|
|
| 82 |
+
| GloVe (6B, 300d) | Static Embedding | ~120M | ~0.50 | 0.0042 |
|
| 83 |
+
| BERT-base (Mean Pool) | Contextual (no NLI FT) | 110M | ~0.50 | 0.0045 |
|
| 84 |
+
| **Collins-Hash (Ours)** | **Hash + RoPE** | **3M** | **0.7114** | **0.237** |
|
| 85 |
+
| paraphrase-MiniLM-L3-v2 | Contextual | 15M | ~0.75 | 0.050 |
|
| 86 |
+
| BGE-micro-v2 | Contextual | 17M | ~0.76 | 0.044 |
|
| 87 |
+
| paraphrase-MiniLM-L6-v2 | Contextual | 22M | ~0.79 | 0.036 |
|
| 88 |
+
| all-mpnet-base-v2 | Contextual | 110M | ~0.83 | 0.0075 |
|
| 89 |
+
|
| 90 |
+
Collins achieves **0.237 score/M** — **5× more efficient** than the next best lightweight model (MiniLM-L3 at 0.050/M), and **53× more efficient** than BERT-base.
|
| 91 |
+
|
| 92 |
+
### Key Findings
|
| 93 |
+
|
| 94 |
+
- **Cross-tier performance**: At 3M params, Collins matches 11–17M parameter models on STSBenchmark — 1/5 the parameters for equivalent semantic fidelity.
|
| 95 |
+
- **Hash compression victory**: MiniLM and ALBERT still carry a full vocabulary embedding table as their largest single component. Collins eliminates this entirely via 2-Universal Hashing.
|
| 96 |
+
- **Sign hash robustness**: STS12–14 scores hold at 0.60–0.62 across diverse domains (news, forums, image captions), confirming differential interference resistance at collision points.
|
| 97 |
+
- **RoPE structural encoding**: STSBenchmark (0.71) > STS12-14 (0.60–0.62) gap indicates stronger performance on well-formed, contextually balanced sentence pairs — exactly where RoPE's topological structure contributes most.
|
| 98 |
+
|
| 99 |
+
---
|
| 100 |
+
|
| 101 |
+
## Applications (2026)
|
| 102 |
+
|
| 103 |
+
This model is designed for deployment scenarios where memory and latency are hard constraints:
|
| 104 |
+
|
| 105 |
+
- **Edge / embedded devices**: Full model fits in 12MB. Suitable for on-device semantic search on mobile, IoT, and microcontrollers with ML accelerators.
|
| 106 |
+
- **Ultra-high-throughput vector search**: L2-cache residency enables millions of encode calls per second on a single GPU, making it viable as the encoder backbone for billion-scale ANN indexes (FAISS, ScaNN, Milvus).
|
| 107 |
+
- **Real-time RAG pipelines**: Sub-millisecond encoding latency unlocks synchronous retrieval in latency-sensitive LLM inference chains without a separate embedding service.
|
| 108 |
+
- **Privacy-preserving on-device NLP**: No network round-trip required. Encode and search entirely on-device for sensitive document workflows.
|
| 109 |
+
- **Low-power inference**: Power consumption scales with model size. At 3M params, Collins is viable on NPU/TPU edge chips where 100M+ models are cost-prohibitive.
|
| 110 |
+
|
| 111 |
+
---
|
| 112 |
+
|
| 113 |
+
## Training
|
| 114 |
+
|
| 115 |
+
- Dataset: `sentence-transformers/all-nli`, triplet split (557,850 samples)
|
| 116 |
+
- Loss: `MultipleNegativesRankingLoss`
|
| 117 |
+
- Epochs: 2, batch size: 256, lr: 2e-4 (cosine schedule), bf16
|
| 118 |
+
|
| 119 |
+
```bash
|
| 120 |
+
python train.py
|
| 121 |
+
```
|
| 122 |
+
|
| 123 |
+
---
|
| 124 |
+
|
| 125 |
+
## Citation
|
| 126 |
+
|
| 127 |
+
```bibtex
|
| 128 |
+
@misc{collins-embedding-3m-2026,
|
| 129 |
+
title = {Collins-Embedding-3M: O(1) Hash Encoding for Efficient Sentence Embeddings},
|
| 130 |
+
author = {NoesisLab},
|
| 131 |
+
year = {2026},
|
| 132 |
+
url = {https://huggingface.co/NoesisLab/Collins-Embedding-3M}
|
| 133 |
+
}
|
| 134 |
+
```
|
collins_sts_comparison.pdf
ADDED
|
Binary file (29 kB). View file
|
|
|
collins_sts_comparison.png
ADDED
|
Git LFS Details
|
config_sentence_transformers.json
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_type": "SentenceTransformer",
|
| 3 |
+
"__version__": {
|
| 4 |
+
"sentence_transformers": "5.2.0",
|
| 5 |
+
"transformers": "4.57.1",
|
| 6 |
+
"pytorch": "2.9.1+cu128"
|
| 7 |
+
},
|
| 8 |
+
"prompts": {
|
| 9 |
+
"query": "",
|
| 10 |
+
"document": ""
|
| 11 |
+
},
|
| 12 |
+
"default_prompt_name": null,
|
| 13 |
+
"similarity_fn_name": "cosine"
|
| 14 |
+
}
|
modeling_hf.py
ADDED
|
@@ -0,0 +1,212 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Collins-RoPE 极简 Embedding 模型(HuggingFace 原生实现)
|
| 3 |
+
架构:Hash Embedding (2-Universal + Sign Hash) -> RoPE -> Transformer Encoder -> Mean Pooling
|
| 4 |
+
目标参数量:~2M
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import math
|
| 8 |
+
from dataclasses import dataclass
|
| 9 |
+
from typing import Optional
|
| 10 |
+
|
| 11 |
+
import torch
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
import torch.nn.functional as F
|
| 14 |
+
from transformers import PretrainedConfig, PreTrainedModel
|
| 15 |
+
from transformers.modeling_outputs import BaseModelOutput
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class CollinsConfig(PretrainedConfig):
|
| 19 |
+
model_type = "collins"
|
| 20 |
+
|
| 21 |
+
def __init__(
|
| 22 |
+
self,
|
| 23 |
+
vocab_size: int = 30522,
|
| 24 |
+
num_buckets: int = 2048,
|
| 25 |
+
hidden_size: int = 256,
|
| 26 |
+
num_hidden_layers: int = 3,
|
| 27 |
+
num_attention_heads: int = 8,
|
| 28 |
+
intermediate_size: int = 1024,
|
| 29 |
+
hidden_dropout_prob: float = 0.1,
|
| 30 |
+
attention_probs_dropout_prob: float = 0.1,
|
| 31 |
+
max_position_embeddings: int = 512,
|
| 32 |
+
# 2-Universal Hash 固定种子(保证 load 后哈希一致)
|
| 33 |
+
hash_seed: int = 42,
|
| 34 |
+
**kwargs,
|
| 35 |
+
):
|
| 36 |
+
super().__init__(**kwargs)
|
| 37 |
+
self.vocab_size = vocab_size
|
| 38 |
+
self.num_buckets = num_buckets
|
| 39 |
+
self.hidden_size = hidden_size
|
| 40 |
+
self.num_hidden_layers = num_hidden_layers
|
| 41 |
+
self.num_attention_heads = num_attention_heads
|
| 42 |
+
self.intermediate_size = intermediate_size
|
| 43 |
+
self.hidden_dropout_prob = hidden_dropout_prob
|
| 44 |
+
self.attention_probs_dropout_prob = attention_probs_dropout_prob
|
| 45 |
+
self.max_position_embeddings = max_position_embeddings
|
| 46 |
+
self.hash_seed = hash_seed
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
class CollinsHashEmbedding(nn.Module):
|
| 50 |
+
"""
|
| 51 |
+
2-Universal Hash + Sign Hash 压缩 Embedding。
|
| 52 |
+
哈希参数从 config.hash_seed 确定性生成,保证 save/load 后一致。
|
| 53 |
+
"""
|
| 54 |
+
|
| 55 |
+
def __init__(self, config: CollinsConfig):
|
| 56 |
+
super().__init__()
|
| 57 |
+
self.num_buckets = config.num_buckets
|
| 58 |
+
self.hidden_size = config.hidden_size
|
| 59 |
+
|
| 60 |
+
self.hash_table = nn.Parameter(
|
| 61 |
+
torch.randn(config.num_buckets, config.hidden_size)
|
| 62 |
+
/ math.sqrt(config.hidden_size)
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
prime = 2147483647 # 梅森素数 2^31 - 1
|
| 66 |
+
rng = torch.Generator()
|
| 67 |
+
rng.manual_seed(config.hash_seed)
|
| 68 |
+
a1 = torch.randint(1, prime, (1,), generator=rng, dtype=torch.long)
|
| 69 |
+
b1 = torch.randint(0, prime, (1,), generator=rng, dtype=torch.long)
|
| 70 |
+
a2 = torch.randint(1, prime, (1,), generator=rng, dtype=torch.long)
|
| 71 |
+
b2 = torch.randint(0, prime, (1,), generator=rng, dtype=torch.long)
|
| 72 |
+
|
| 73 |
+
self.register_buffer("prime", torch.tensor(prime, dtype=torch.long))
|
| 74 |
+
self.register_buffer("a1", a1)
|
| 75 |
+
self.register_buffer("b1", b1)
|
| 76 |
+
self.register_buffer("a2", a2)
|
| 77 |
+
self.register_buffer("b2", b2)
|
| 78 |
+
|
| 79 |
+
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
|
| 80 |
+
x = input_ids.long()
|
| 81 |
+
bucket_idx = ((x * self.a1 + self.b1) % self.prime) % self.num_buckets
|
| 82 |
+
sign = ((x * self.a2 + self.b2) % self.prime) % 2
|
| 83 |
+
sign = (sign * 2 - 1).float()
|
| 84 |
+
return self.hash_table[bucket_idx] * sign.unsqueeze(-1)
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
class CollinsModel(PreTrainedModel):
|
| 88 |
+
"""
|
| 89 |
+
Collins-RoPE Encoder,输出 last_hidden_state 和 pooler_output。
|
| 90 |
+
使用 transformers.models.bert 的 BertEncoder + RoPE 替换 BertEmbeddings。
|
| 91 |
+
"""
|
| 92 |
+
|
| 93 |
+
config_class = CollinsConfig
|
| 94 |
+
base_model_prefix = "collins"
|
| 95 |
+
supports_gradient_checkpointing = True
|
| 96 |
+
|
| 97 |
+
def __init__(self, config: CollinsConfig):
|
| 98 |
+
super().__init__(config)
|
| 99 |
+
self.config = config
|
| 100 |
+
|
| 101 |
+
self.embeddings = CollinsHashEmbedding(config)
|
| 102 |
+
|
| 103 |
+
# 直接复用 HF BertEncoder(含 Multi-Head Attention + FFN + LayerNorm)
|
| 104 |
+
from transformers.models.bert.modeling_bert import BertEncoder, BertConfig
|
| 105 |
+
|
| 106 |
+
bert_cfg = BertConfig(
|
| 107 |
+
hidden_size=config.hidden_size,
|
| 108 |
+
num_hidden_layers=config.num_hidden_layers,
|
| 109 |
+
num_attention_heads=config.num_attention_heads,
|
| 110 |
+
intermediate_size=config.intermediate_size,
|
| 111 |
+
hidden_dropout_prob=config.hidden_dropout_prob,
|
| 112 |
+
attention_probs_dropout_prob=config.attention_probs_dropout_prob,
|
| 113 |
+
max_position_embeddings=config.max_position_embeddings,
|
| 114 |
+
# 关闭 Bert 自带的位置编码,我们用 RoPE
|
| 115 |
+
position_embedding_type="relative_key_query",
|
| 116 |
+
)
|
| 117 |
+
bert_cfg._attn_implementation = "eager"
|
| 118 |
+
self.encoder = BertEncoder(bert_cfg)
|
| 119 |
+
|
| 120 |
+
# RoPE 频率缓冲(无参数)
|
| 121 |
+
dim = config.hidden_size
|
| 122 |
+
inv_freq = 1.0 / (
|
| 123 |
+
10000 ** (torch.arange(0, dim, 2).float() / dim)
|
| 124 |
+
)
|
| 125 |
+
t = torch.arange(config.max_position_embeddings).float()
|
| 126 |
+
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
| 127 |
+
self.register_buffer("rope_cos", freqs.cos())
|
| 128 |
+
self.register_buffer("rope_sin", freqs.sin())
|
| 129 |
+
|
| 130 |
+
self.post_init()
|
| 131 |
+
|
| 132 |
+
def _apply_rope(self, x: torch.Tensor) -> torch.Tensor:
|
| 133 |
+
seq_len = x.shape[1]
|
| 134 |
+
cos = self.rope_cos[:seq_len].unsqueeze(0)
|
| 135 |
+
sin = self.rope_sin[:seq_len].unsqueeze(0)
|
| 136 |
+
x1, x2 = x[..., 0::2], x[..., 1::2]
|
| 137 |
+
return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
|
| 138 |
+
|
| 139 |
+
def get_extended_attention_mask(self, attention_mask: torch.Tensor) -> torch.Tensor:
|
| 140 |
+
# BertEncoder 需要 [B, 1, 1, L] 形式的 mask,0 = 保留,-inf = 忽略
|
| 141 |
+
extended = attention_mask[:, None, None, :]
|
| 142 |
+
extended = (1.0 - extended.float()) * torch.finfo(torch.float32).min
|
| 143 |
+
return extended
|
| 144 |
+
|
| 145 |
+
def forward(
|
| 146 |
+
self,
|
| 147 |
+
input_ids: torch.Tensor,
|
| 148 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 149 |
+
return_dict: bool = True,
|
| 150 |
+
):
|
| 151 |
+
if attention_mask is None:
|
| 152 |
+
attention_mask = torch.ones_like(input_ids)
|
| 153 |
+
|
| 154 |
+
x = self.embeddings(input_ids) # [B, L, D]
|
| 155 |
+
x = self._apply_rope(x) # [B, L, D]
|
| 156 |
+
|
| 157 |
+
ext_mask = self.get_extended_attention_mask(attention_mask)
|
| 158 |
+
encoder_out = self.encoder(x, attention_mask=ext_mask)
|
| 159 |
+
hidden_states = encoder_out.last_hidden_state # [B, L, D]
|
| 160 |
+
|
| 161 |
+
# Mean Pooling
|
| 162 |
+
mask = attention_mask.unsqueeze(-1).float()
|
| 163 |
+
pooled = (hidden_states * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
|
| 164 |
+
pooled = F.normalize(pooled, p=2, dim=-1)
|
| 165 |
+
|
| 166 |
+
if not return_dict:
|
| 167 |
+
return (hidden_states, pooled)
|
| 168 |
+
|
| 169 |
+
return BaseModelOutput(
|
| 170 |
+
last_hidden_state=hidden_states,
|
| 171 |
+
hidden_states=None,
|
| 172 |
+
attentions=None,
|
| 173 |
+
), pooled
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
class CollinsSTWrapper(nn.Module):
|
| 177 |
+
"""
|
| 178 |
+
sentence-transformers 5.x 兼容包装层。
|
| 179 |
+
持有 tokenizer,实现 tokenize() 接口,同时注入 sentence_embedding。
|
| 180 |
+
"""
|
| 181 |
+
|
| 182 |
+
def __init__(self, collins_model: CollinsModel, tokenizer_name_or_path: str = "bert-base-uncased", max_seq_length: int = 128):
|
| 183 |
+
super().__init__()
|
| 184 |
+
from transformers import AutoTokenizer
|
| 185 |
+
self.collins_model = collins_model
|
| 186 |
+
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
|
| 187 |
+
self.max_seq_length = max_seq_length
|
| 188 |
+
|
| 189 |
+
def tokenize(self, texts: list[str], padding: str | bool = True) -> dict:
|
| 190 |
+
return self.tokenizer(
|
| 191 |
+
texts,
|
| 192 |
+
padding=padding,
|
| 193 |
+
truncation=True,
|
| 194 |
+
max_length=self.max_seq_length,
|
| 195 |
+
return_tensors="pt",
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
def forward(self, features: dict) -> dict:
|
| 199 |
+
input_ids = features["input_ids"]
|
| 200 |
+
attention_mask = features.get("attention_mask", None)
|
| 201 |
+
_, pooled = self.collins_model(input_ids, attention_mask)
|
| 202 |
+
features["sentence_embedding"] = pooled
|
| 203 |
+
return features
|
| 204 |
+
|
| 205 |
+
def save(self, output_path: str):
|
| 206 |
+
self.collins_model.save_pretrained(output_path)
|
| 207 |
+
self.tokenizer.save_pretrained(output_path)
|
| 208 |
+
|
| 209 |
+
@staticmethod
|
| 210 |
+
def load(input_path: str) -> "CollinsSTWrapper":
|
| 211 |
+
model = CollinsModel.from_pretrained(input_path)
|
| 212 |
+
return CollinsSTWrapper(model, tokenizer_name_or_path=input_path)
|
modules.json
ADDED
|
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"idx": 0,
|
| 4 |
+
"name": "0",
|
| 5 |
+
"path": "0_CollinsSTWrapper",
|
| 6 |
+
"type": "modeling_hf.CollinsSTWrapper"
|
| 7 |
+
}
|
| 8 |
+
]
|