Any-to-Any
Transformers
Safetensors
multilingual
spev
feature-extraction
multimodal
text
audio
image
video
decision-model
choice
qwen3-omni
lora
custom-code
custom_code
Instructions to use SpragAI/SPEV with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use SpragAI/SPEV with Transformers:
# pip install -U transformers accelerate # Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("SpragAI/SPEV", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
Release isolated single-pass SPEV decision head
Browse files- MANIFEST.json +0 -164
- README.md +2 -13
- __init__.py +0 -0
- branch_attention.py +115 -0
- candidate_tree.py +52 -0
- config.json +5 -1
- configuration_spev.py +15 -0
- decision_head.py +126 -0
- decision_layout.py +84 -0
- decision_utils.py +1 -1
- examples/turn_decision.py +0 -28
- examples/turn_policy.json +0 -17
- media.py +50 -0
- model-decision.safetensors +2 -2
- model.safetensors.index.json +85 -5
- modeling_spev.py +65 -48
- processing_spev.py +26 -17
- processor_config.json +3 -0
- token_operations.py +78 -0
MANIFEST.json
DELETED
|
@@ -1,164 +0,0 @@
|
|
| 1 |
-
{
|
| 2 |
-
"model": "SpragAI/SPEV",
|
| 3 |
-
"package_format": "transformers-autoclasses",
|
| 4 |
-
"candidate": "clean57/joint/candidate",
|
| 5 |
-
"source_adapter_sha256": "1e1ffd17099c9e70ab061547286a43a6bdddd8ce8aa1d304803f65550fc8a7b5",
|
| 6 |
-
"source_head_sha256": "6f1c6891aec8cb6862493ce24c3895e33e2c5b95ce8781c0ab61dd27b0702f94",
|
| 7 |
-
"files": [
|
| 8 |
-
{
|
| 9 |
-
"path": "LICENSE",
|
| 10 |
-
"bytes": 11358,
|
| 11 |
-
"sha256": "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"
|
| 12 |
-
},
|
| 13 |
-
{
|
| 14 |
-
"path": "README.md",
|
| 15 |
-
"bytes": 4119,
|
| 16 |
-
"sha256": "73d6815884bd212924e2cd7450f42897ac67bc4feeb9883eaa9d8d39bb287993"
|
| 17 |
-
},
|
| 18 |
-
{
|
| 19 |
-
"path": "THIRD_PARTY_NOTICES.md",
|
| 20 |
-
"bytes": 358,
|
| 21 |
-
"sha256": "4c8a4cecaf46c00fc31b21aaed23e9c15099b312f07e922ecbf49945df683b45"
|
| 22 |
-
},
|
| 23 |
-
{
|
| 24 |
-
"path": "__init__.py",
|
| 25 |
-
"bytes": 0,
|
| 26 |
-
"sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
|
| 27 |
-
},
|
| 28 |
-
{
|
| 29 |
-
"path": "chat_template.jinja",
|
| 30 |
-
"bytes": 6519,
|
| 31 |
-
"sha256": "45777d2e798aae1d1ce6844e69edcb23fc0c710009554711c78efab9044da2c8"
|
| 32 |
-
},
|
| 33 |
-
{
|
| 34 |
-
"path": "config.json",
|
| 35 |
-
"bytes": 6341,
|
| 36 |
-
"sha256": "98f5092ed524600d44113f192c45c2ef8cb078d4ae77885ab2376db88a818b4f"
|
| 37 |
-
},
|
| 38 |
-
{
|
| 39 |
-
"path": "configuration_spev.py",
|
| 40 |
-
"bytes": 1750,
|
| 41 |
-
"sha256": "8390f4ded6259839601d878e62c4f44559b44999f63c3ae141879a911003de03"
|
| 42 |
-
},
|
| 43 |
-
{
|
| 44 |
-
"path": "decision_utils.py",
|
| 45 |
-
"bytes": 3867,
|
| 46 |
-
"sha256": "bc38e41c1123880b2038696489272cb489a1c62da98f672f467e694078636208"
|
| 47 |
-
},
|
| 48 |
-
{
|
| 49 |
-
"path": "examples/turn_decision.py",
|
| 50 |
-
"bytes": 981,
|
| 51 |
-
"sha256": "e1ce54067b9cd8b6348b97309ee7920c6a0eadddf75880676f3335ae45c883d8"
|
| 52 |
-
},
|
| 53 |
-
{
|
| 54 |
-
"path": "examples/turn_policy.json",
|
| 55 |
-
"bytes": 730,
|
| 56 |
-
"sha256": "d07af2090edf412ab86af84887b11e7d96e09b22325799788b84921d130e3406"
|
| 57 |
-
},
|
| 58 |
-
{
|
| 59 |
-
"path": "model-00001-of-00013.safetensors",
|
| 60 |
-
"bytes": 4375707904,
|
| 61 |
-
"sha256": "27a8786a74a21747745003b524f343d7d39b5acfe3edf4ce8127cbfc7ccc7788"
|
| 62 |
-
},
|
| 63 |
-
{
|
| 64 |
-
"path": "model-00002-of-00013.safetensors",
|
| 65 |
-
"bytes": 4985172432,
|
| 66 |
-
"sha256": "ec6b95123dead3f88cc7b7ee87da4ab49cb50878c7ae7446d0b8d2b81990fe13"
|
| 67 |
-
},
|
| 68 |
-
{
|
| 69 |
-
"path": "model-00003-of-00013.safetensors",
|
| 70 |
-
"bytes": 4985172432,
|
| 71 |
-
"sha256": "14d8680d304fcedd57440d3722591ad3347530d2e1d3952713cdbffcbc4bce00"
|
| 72 |
-
},
|
| 73 |
-
{
|
| 74 |
-
"path": "model-00004-of-00013.safetensors",
|
| 75 |
-
"bytes": 4985173960,
|
| 76 |
-
"sha256": "82223f49e7ad8eb263f442642f6b260bb251df8a7703968e36289b04845cf228"
|
| 77 |
-
},
|
| 78 |
-
{
|
| 79 |
-
"path": "model-00005-of-00013.safetensors",
|
| 80 |
-
"bytes": 4985174008,
|
| 81 |
-
"sha256": "ff1cd35f862c04aa8bbb8cd2c7ce909c8833866964170185fe6ef69ee6f79328"
|
| 82 |
-
},
|
| 83 |
-
{
|
| 84 |
-
"path": "model-00006-of-00013.safetensors",
|
| 85 |
-
"bytes": 4985174008,
|
| 86 |
-
"sha256": "efb7cf9badc49260df295f1f3b9ee852195643579a22b17930ed847cee1b801f"
|
| 87 |
-
},
|
| 88 |
-
{
|
| 89 |
-
"path": "model-00007-of-00013.safetensors",
|
| 90 |
-
"bytes": 4985174008,
|
| 91 |
-
"sha256": "aeeed3d751bd7f4b8a4d49c56fd4f84ce5ca6352e989136ac0f7be9f7a1266ad"
|
| 92 |
-
},
|
| 93 |
-
{
|
| 94 |
-
"path": "model-00008-of-00013.safetensors",
|
| 95 |
-
"bytes": 4985174008,
|
| 96 |
-
"sha256": "1b1ec719c4854406189f4727cd57eaf6f29756eae7a5a02afaaa351c45fecde3"
|
| 97 |
-
},
|
| 98 |
-
{
|
| 99 |
-
"path": "model-00009-of-00013.safetensors",
|
| 100 |
-
"bytes": 4985174008,
|
| 101 |
-
"sha256": "b85cf0be2c40324cb3372de572b6e38e60e02264d9444eb0c2054aac98984d60"
|
| 102 |
-
},
|
| 103 |
-
{
|
| 104 |
-
"path": "model-00010-of-00013.safetensors",
|
| 105 |
-
"bytes": 4985174008,
|
| 106 |
-
"sha256": "4cae94eb7f02d03a1ae90d055b6119ae8b3fdc877827ee714ec41f33029007cf"
|
| 107 |
-
},
|
| 108 |
-
{
|
| 109 |
-
"path": "model-00011-of-00013.safetensors",
|
| 110 |
-
"bytes": 4985174008,
|
| 111 |
-
"sha256": "9b74e898729090f19b9e0ba76a9f4736bd6fcc09ebb1422d1717a9b95645ece7"
|
| 112 |
-
},
|
| 113 |
-
{
|
| 114 |
-
"path": "model-00012-of-00013.safetensors",
|
| 115 |
-
"bytes": 4985174008,
|
| 116 |
-
"sha256": "6d6a2d55237097a1086f1522fb41fef8877873f04a110a4393e4e218fd28726f"
|
| 117 |
-
},
|
| 118 |
-
{
|
| 119 |
-
"path": "model-00013-of-00013.safetensors",
|
| 120 |
-
"bytes": 3605546880,
|
| 121 |
-
"sha256": "089e485dcc0bb546286a7f0acc906d15f080502d7bd9ee4b2441b9e2bb33be41"
|
| 122 |
-
},
|
| 123 |
-
{
|
| 124 |
-
"path": "model-decision.safetensors",
|
| 125 |
-
"bytes": 55615444,
|
| 126 |
-
"sha256": "e711007c2b401a337442e7b9e1ae6d89c36a2a2dee67cda9a817730b766ca1a1"
|
| 127 |
-
},
|
| 128 |
-
{
|
| 129 |
-
"path": "model.safetensors.index.json",
|
| 130 |
-
"bytes": 1964764,
|
| 131 |
-
"sha256": "61c7968601e201e1713058cc0a32619ce0edd8958b574ce7f442a6c1fa4e44d4"
|
| 132 |
-
},
|
| 133 |
-
{
|
| 134 |
-
"path": "modeling_spev.py",
|
| 135 |
-
"bytes": 4750,
|
| 136 |
-
"sha256": "625f3c60581eed62a278c23f68cd9494c16023707e9d8ddb30e3ae47fcd6ba87"
|
| 137 |
-
},
|
| 138 |
-
{
|
| 139 |
-
"path": "processing_spev.py",
|
| 140 |
-
"bytes": 4330,
|
| 141 |
-
"sha256": "821524b54c2de7a7cc6800dc8fe4679318542a86cdbffc6646fe82ccc64e1a96"
|
| 142 |
-
},
|
| 143 |
-
{
|
| 144 |
-
"path": "processor_config.json",
|
| 145 |
-
"bytes": 2704,
|
| 146 |
-
"sha256": "5712ee707a83a1c8fbc230df116395e6c61be58770f648a61cbf8c504f9dc380"
|
| 147 |
-
},
|
| 148 |
-
{
|
| 149 |
-
"path": "requirements.txt",
|
| 150 |
-
"bytes": 205,
|
| 151 |
-
"sha256": "4b107e5fca893565799ebe5b2600c5708d28d618c0d581125bdfa3df52c15435"
|
| 152 |
-
},
|
| 153 |
-
{
|
| 154 |
-
"path": "tokenizer.json",
|
| 155 |
-
"bytes": 11424262,
|
| 156 |
-
"sha256": "a7d41145c408f3062e824f965be1c29854cd809e111cdf8170aa8b0bcd5d1fab"
|
| 157 |
-
},
|
| 158 |
-
{
|
| 159 |
-
"path": "tokenizer_config.json",
|
| 160 |
-
"bytes": 1067,
|
| 161 |
-
"sha256": "6d6cadf5fc644eef4fe90b62f053f759733e4543921f52f37adb22cc48492779"
|
| 162 |
-
}
|
| 163 |
-
]
|
| 164 |
-
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
README.md
CHANGED
|
@@ -29,7 +29,7 @@ returns an answer with a probability for each choice.
|
|
| 29 |
## Model
|
| 30 |
|
| 31 |
- **Backbone:** Qwen3-Omni-30B-A3B Thinker, with a trained LoRA adapter.
|
| 32 |
-
- **Decision head:**
|
| 33 |
- **Output:** per-choice probabilities and the highest-probability answer for each question.
|
| 34 |
|
| 35 |
## Quickstart
|
|
@@ -104,23 +104,12 @@ do not supply a sample-rate argument. Supply a video's soundtrack separately as
|
|
| 104 |
| `model-*.safetensors`, `model.safetensors.index.json` | Backbone, adapter and decision-head weights |
|
| 105 |
| `config.json`, `configuration_spev.py` | Model architecture and scoring configuration |
|
| 106 |
| `modeling_spev.py` | Transformers model implementation |
|
|
|
|
| 107 |
| `processing_spev.py`, `decision_utils.py` | Media preparation, question rendering and answer formatting |
|
| 108 |
| Tokenizer, processor and chat-template files | Text, audio, image and video processing |
|
| 109 |
-
| `examples/turn_policy.json` | Evaluated turn-taking question and application threshold |
|
| 110 |
| `requirements.txt` | Inference dependencies |
|
| 111 |
| `LICENSE` | Apache-2.0 license |
|
| 112 |
|
| 113 |
-
## Results
|
| 114 |
-
|
| 115 |
-
| Input | SPEV p50 / p95 | OpenAI p50 / p95 |
|
| 116 |
-
|---|---:|---:|
|
| 117 |
-
| 5-second audio | 166 / 191 ms | 752 / 1331 ms |
|
| 118 |
-
| 20-second audio | 216 / 270 ms | 997 / 1326 ms |
|
| 119 |
-
| 512 × 512 image | 183 / 283 ms | 194 / 334 ms |
|
| 120 |
-
| 1024 × 1024 image | 322 / 642 ms | 373 / 763 ms |
|
| 121 |
-
|
| 122 |
-
Warm public-HTTPS latency including RTT, measured from the same client with concurrency one. SPEV ran on an RTX PRO6000 with native vLLM pooling. Each row has 30 measured pairs over six fixtures after three warmup pairs. OpenAI used `gpt-6-luna`; audio includes a fresh `gpt-4o-mini-transcribe` request before Decisions. These measurements cover latency, not accuracy.
|
| 123 |
-
|
| 124 |
## License
|
| 125 |
|
| 126 |
Apache-2.0. See [LICENSE](LICENSE) and [third-party notices](THIRD_PARTY_NOTICES.md).
|
|
|
|
| 29 |
## Model
|
| 30 |
|
| 31 |
- **Backbone:** Qwen3-Omni-30B-A3B Thinker, with a trained LoRA adapter.
|
| 32 |
+
- **Decision head:** an evidence-attention head that scores each question independently.
|
| 33 |
- **Output:** per-choice probabilities and the highest-probability answer for each question.
|
| 34 |
|
| 35 |
## Quickstart
|
|
|
|
| 104 |
| `model-*.safetensors`, `model.safetensors.index.json` | Backbone, adapter and decision-head weights |
|
| 105 |
| `config.json`, `configuration_spev.py` | Model architecture and scoring configuration |
|
| 106 |
| `modeling_spev.py` | Transformers model implementation |
|
| 107 |
+
| `decision_head.py`, `decision_layout.py`, `candidate_tree.py`, `branch_attention.py`, `token_operations.py` | Decision scoring and question isolation |
|
| 108 |
| `processing_spev.py`, `decision_utils.py` | Media preparation, question rendering and answer formatting |
|
| 109 |
| Tokenizer, processor and chat-template files | Text, audio, image and video processing |
|
|
|
|
| 110 |
| `requirements.txt` | Inference dependencies |
|
| 111 |
| `LICENSE` | Apache-2.0 license |
|
| 112 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 113 |
## License
|
| 114 |
|
| 115 |
Apache-2.0. See [LICENSE](LICENSE) and [third-party notices](THIRD_PARTY_NOTICES.md).
|
__init__.py
DELETED
|
File without changes
|
branch_attention.py
ADDED
|
@@ -0,0 +1,115 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
| 3 |
+
"""Shared-prefix branch attention with request-local immutable mask metadata."""
|
| 4 |
+
|
| 5 |
+
from collections import defaultdict
|
| 6 |
+
from contextlib import contextmanager
|
| 7 |
+
from contextvars import ContextVar
|
| 8 |
+
|
| 9 |
+
import torch
|
| 10 |
+
from torch.nn import functional as F
|
| 11 |
+
|
| 12 |
+
REQUEST = ContextVar("spev_attention_request", default=None)
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class RequestMaskCache:
|
| 16 |
+
def __init__(self):
|
| 17 |
+
self.tree = None
|
| 18 |
+
self.device = None
|
| 19 |
+
self.plan = None
|
| 20 |
+
self.builds = 0
|
| 21 |
+
self.hits = 0
|
| 22 |
+
|
| 23 |
+
def __call__(self, tree, device):
|
| 24 |
+
if tree is self.tree and device == self.device:
|
| 25 |
+
self.hits += 1
|
| 26 |
+
return self.plan
|
| 27 |
+
buckets = defaultdict(list)
|
| 28 |
+
start = tree.prefix_length
|
| 29 |
+
for length in tree.branch_lengths:
|
| 30 |
+
buckets[((length + 31) // 32) * 32].append((start, length))
|
| 31 |
+
start += length
|
| 32 |
+
assert start == tree.length
|
| 33 |
+
plan = []
|
| 34 |
+
p = tree.prefix_length
|
| 35 |
+
for width, branches in sorted(buckets.items()):
|
| 36 |
+
for offset in range(0, len(branches), 8):
|
| 37 |
+
block = branches[offset : offset + 8]
|
| 38 |
+
allowed = torch.zeros(8, 1, width, p + width, device=device, dtype=torch.bool)
|
| 39 |
+
allowed[:, :, :, 0] = True
|
| 40 |
+
for slot, (_, length) in enumerate(block):
|
| 41 |
+
allowed[slot, 0, :length, :p] = True
|
| 42 |
+
allowed[slot, 0, :length, p : p + length] = torch.ones(
|
| 43 |
+
length, length, device=device, dtype=torch.bool
|
| 44 |
+
).tril()
|
| 45 |
+
plan.append((width, block, allowed))
|
| 46 |
+
self.tree = tree
|
| 47 |
+
self.device = device
|
| 48 |
+
self.plan = plan
|
| 49 |
+
self.builds += 1
|
| 50 |
+
return plan
|
| 51 |
+
|
| 52 |
+
def stats(self):
|
| 53 |
+
return dict(
|
| 54 |
+
builds=self.builds,
|
| 55 |
+
hits=self.hits,
|
| 56 |
+
blocks=len(self.plan or []),
|
| 57 |
+
branches=sum(len(block) for _, block, _ in self.plan or []),
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def branch_attention(
|
| 62 |
+
module, query, key, value, attention_mask=None, dropout=0.0, scaling=None, sliding_window=None, **kwargs
|
| 63 |
+
):
|
| 64 |
+
request = REQUEST.get()
|
| 65 |
+
tree = request[0] if request else None
|
| 66 |
+
if tree is None or query.shape[0] != 1 or key.shape[-2] != tree.length or query.shape[-2] != tree.length:
|
| 67 |
+
raise ValueError("Branch attention requires an explicit single-request tree layout")
|
| 68 |
+
if sliding_window is not None:
|
| 69 |
+
raise ValueError("Sliding-window branch attention is not implemented")
|
| 70 |
+
if dropout:
|
| 71 |
+
raise ValueError("Branch attention reference requires zero attention dropout")
|
| 72 |
+
groups = query.shape[1] // key.shape[1]
|
| 73 |
+
if groups > 1:
|
| 74 |
+
key = key.repeat_interleave(groups, dim=1)
|
| 75 |
+
value = value.repeat_interleave(groups, dim=1)
|
| 76 |
+
p = tree.prefix_length
|
| 77 |
+
output = torch.zeros_like(query)
|
| 78 |
+
output[:, :, :p] = F.scaled_dot_product_attention(
|
| 79 |
+
query[:, :, :p], key[:, :, :p], value[:, :, :p], is_causal=True, scale=scaling
|
| 80 |
+
)
|
| 81 |
+
cache = request[1]
|
| 82 |
+
heads, dim = query.shape[1], query.shape[-1]
|
| 83 |
+
batch_slots = 8
|
| 84 |
+
for width, block, allowed in cache(tree, query.device):
|
| 85 |
+
q = query.new_zeros(batch_slots, heads, width, dim)
|
| 86 |
+
k = key.new_zeros(batch_slots, heads, p + width, dim)
|
| 87 |
+
v = value.new_zeros(batch_slots, heads, p + width, dim)
|
| 88 |
+
for slot, (start, length) in enumerate(block):
|
| 89 |
+
q[slot, :, :length] = query[0, :, start : start + length]
|
| 90 |
+
k[slot, :, :p] = key[0, :, :p]
|
| 91 |
+
v[slot, :, :p] = value[0, :, :p]
|
| 92 |
+
k[slot, :, p : p + length] = key[0, :, start : start + length]
|
| 93 |
+
v[slot, :, p : p + length] = value[0, :, start : start + length]
|
| 94 |
+
values = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, scale=scaling)
|
| 95 |
+
for slot, (start, length) in enumerate(block):
|
| 96 |
+
output[0, :, start : start + length] = values[slot, :, :length]
|
| 97 |
+
return output.transpose(1, 2).contiguous(), None
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def install(model):
|
| 101 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
|
| 102 |
+
|
| 103 |
+
implementation = f"{__name__}.isolated_branches"
|
| 104 |
+
ALL_ATTENTION_FUNCTIONS.register(implementation, branch_attention)
|
| 105 |
+
for layer in model.model.layers:
|
| 106 |
+
layer.self_attn.config._attn_implementation = implementation
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
@contextmanager
|
| 110 |
+
def request_context(tree):
|
| 111 |
+
token = REQUEST.set((tree, RequestMaskCache()))
|
| 112 |
+
try:
|
| 113 |
+
yield
|
| 114 |
+
finally:
|
| 115 |
+
REQUEST.reset(token)
|
candidate_tree.py
ADDED
|
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
| 3 |
+
"""Shared evidence prefix and isolated candidate branch layout."""
|
| 4 |
+
|
| 5 |
+
from dataclasses import dataclass
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
@dataclass(frozen=True)
|
| 11 |
+
class CandidateTree:
|
| 12 |
+
prefix_length: int
|
| 13 |
+
branch_lengths: tuple[int, ...]
|
| 14 |
+
|
| 15 |
+
def __post_init__(self):
|
| 16 |
+
if self.prefix_length <= 0 or not self.branch_lengths or min(self.branch_lengths) <= 0:
|
| 17 |
+
raise ValueError("A nonempty prefix and nonempty branches are required")
|
| 18 |
+
|
| 19 |
+
@property
|
| 20 |
+
def length(self):
|
| 21 |
+
return self.prefix_length + sum(self.branch_lengths)
|
| 22 |
+
|
| 23 |
+
def allowed(self, *, device=None):
|
| 24 |
+
groups = [-1] * self.prefix_length
|
| 25 |
+
for index, length in enumerate(self.branch_lengths):
|
| 26 |
+
groups.extend([index] * length)
|
| 27 |
+
groups = torch.tensor(groups, device=device)
|
| 28 |
+
physical = torch.arange(self.length, device=device)
|
| 29 |
+
causal = physical[None, :] <= physical[:, None]
|
| 30 |
+
return causal & ((groups[None, :] == -1) | (groups[:, None] == groups[None, :]))
|
| 31 |
+
|
| 32 |
+
def additive_mask(self, *, dtype=torch.float32, device=None):
|
| 33 |
+
allowed = self.allowed(device=device)
|
| 34 |
+
return torch.zeros_like(allowed, dtype=dtype).masked_fill(~allowed, float("-inf"))[None, None]
|
| 35 |
+
|
| 36 |
+
def position_ids(self, prefix_positions):
|
| 37 |
+
if prefix_positions.ndim != 3 or prefix_positions.shape[-1] != self.prefix_length:
|
| 38 |
+
raise ValueError("Expected native prefix positions with shape (axes, batch, prefix)")
|
| 39 |
+
if prefix_positions.shape[1] != 1:
|
| 40 |
+
raise ValueError("This reference packs one request")
|
| 41 |
+
origin = int(prefix_positions.max()) + 1
|
| 42 |
+
parts = [prefix_positions]
|
| 43 |
+
for length in self.branch_lengths:
|
| 44 |
+
branch = torch.arange(origin, origin + length, device=prefix_positions.device)
|
| 45 |
+
parts.append(branch.view(1, 1, -1).expand(prefix_positions.shape[0], 1, -1))
|
| 46 |
+
return torch.cat(parts, dim=-1)
|
| 47 |
+
|
| 48 |
+
def attention_pairs(self):
|
| 49 |
+
prefix = self.prefix_length
|
| 50 |
+
return prefix * (prefix + 1) // 2 + sum(
|
| 51 |
+
prefix * length + length * (length + 1) // 2 for length in self.branch_lengths
|
| 52 |
+
)
|
config.json
CHANGED
|
@@ -9,10 +9,14 @@
|
|
| 9 |
"AutoModel": "modeling_spev.SPEVModel"
|
| 10 |
},
|
| 11 |
"dtype": "bfloat16",
|
|
|
|
|
|
|
|
|
|
|
|
|
| 12 |
"head_width": 256,
|
| 13 |
"max_tokens": 8192,
|
| 14 |
"model_type": "spev",
|
| 15 |
-
"temperature": 1.
|
| 16 |
"thinker_config": {
|
| 17 |
"architectures": [
|
| 18 |
"Qwen3OmniMoeThinkerForConditionalGeneration"
|
|
|
|
| 9 |
"AutoModel": "modeling_spev.SPEVModel"
|
| 10 |
},
|
| 11 |
"dtype": "bfloat16",
|
| 12 |
+
"head_architecture": "isolated_evidence_v1",
|
| 13 |
+
"head_depth": 2,
|
| 14 |
+
"head_ff": 1024,
|
| 15 |
+
"head_heads": 8,
|
| 16 |
"head_width": 256,
|
| 17 |
"max_tokens": 8192,
|
| 18 |
"model_type": "spev",
|
| 19 |
+
"temperature": 1.0,
|
| 20 |
"thinker_config": {
|
| 21 |
"architectures": [
|
| 22 |
"Qwen3OmniMoeThinkerForConditionalGeneration"
|
configuration_spev.py
CHANGED
|
@@ -16,6 +16,10 @@ class SPEVConfig(PretrainedConfig):
|
|
| 16 |
adapter_rank=32,
|
| 17 |
adapter_alpha=64,
|
| 18 |
head_width=256,
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
temperature=1.0,
|
| 20 |
max_tokens=8192,
|
| 21 |
**kwargs,
|
|
@@ -26,6 +30,9 @@ class SPEVConfig(PretrainedConfig):
|
|
| 26 |
for name, value in (
|
| 27 |
("adapter_rank", adapter_rank),
|
| 28 |
("head_width", head_width),
|
|
|
|
|
|
|
|
|
|
| 29 |
("max_tokens", max_tokens),
|
| 30 |
):
|
| 31 |
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
|
@@ -37,6 +44,14 @@ class SPEVConfig(PretrainedConfig):
|
|
| 37 |
self.adapter_rank = adapter_rank
|
| 38 |
self.adapter_alpha = adapter_alpha
|
| 39 |
self.head_width = head_width
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 40 |
self.temperature = temperature
|
| 41 |
self.max_tokens = max_tokens
|
| 42 |
kwargs.setdefault("tie_word_embeddings", False)
|
|
|
|
| 16 |
adapter_rank=32,
|
| 17 |
adapter_alpha=64,
|
| 18 |
head_width=256,
|
| 19 |
+
head_heads=8,
|
| 20 |
+
head_ff=1024,
|
| 21 |
+
head_depth=2,
|
| 22 |
+
head_architecture="isolated_evidence_v1",
|
| 23 |
temperature=1.0,
|
| 24 |
max_tokens=8192,
|
| 25 |
**kwargs,
|
|
|
|
| 30 |
for name, value in (
|
| 31 |
("adapter_rank", adapter_rank),
|
| 32 |
("head_width", head_width),
|
| 33 |
+
("head_heads", head_heads),
|
| 34 |
+
("head_ff", head_ff),
|
| 35 |
+
("head_depth", head_depth),
|
| 36 |
("max_tokens", max_tokens),
|
| 37 |
):
|
| 38 |
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
|
|
|
| 44 |
self.adapter_rank = adapter_rank
|
| 45 |
self.adapter_alpha = adapter_alpha
|
| 46 |
self.head_width = head_width
|
| 47 |
+
if head_architecture != "isolated_evidence_v1":
|
| 48 |
+
raise ValueError("Unsupported SPEV decision head architecture")
|
| 49 |
+
if head_width % head_heads:
|
| 50 |
+
raise ValueError("head_width must be divisible by head_heads")
|
| 51 |
+
self.head_heads = head_heads
|
| 52 |
+
self.head_ff = head_ff
|
| 53 |
+
self.head_depth = head_depth
|
| 54 |
+
self.head_architecture = head_architecture
|
| 55 |
self.temperature = temperature
|
| 56 |
self.max_tokens = max_tokens
|
| 57 |
kwargs.setdefault("tie_word_embeddings", False)
|
decision_head.py
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
| 3 |
+
"""Question-local evidence routing for SPEV decision checkpoints."""
|
| 4 |
+
|
| 5 |
+
from dataclasses import dataclass
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
from torch import nn
|
| 9 |
+
from torch.nn import functional as F
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class EvidenceBlock(nn.Module):
|
| 13 |
+
def __init__(self, width: int, heads: int, ff: int):
|
| 14 |
+
super().__init__()
|
| 15 |
+
self.qnorm = nn.LayerNorm(width)
|
| 16 |
+
self.mnorm = nn.LayerNorm(width)
|
| 17 |
+
self.attn = nn.MultiheadAttention(width, heads, batch_first=True, dropout=0.0)
|
| 18 |
+
self.fnorm = nn.LayerNorm(width)
|
| 19 |
+
self.ff = nn.Sequential(nn.Linear(width, ff), nn.GELU(), nn.Linear(ff, width))
|
| 20 |
+
|
| 21 |
+
def forward(self, query, memory, memory_pad):
|
| 22 |
+
memory = self.mnorm(memory)
|
| 23 |
+
value, _ = self.attn(self.qnorm(query), memory, memory, key_padding_mask=memory_pad, need_weights=False)
|
| 24 |
+
query = query + value
|
| 25 |
+
return query + self.ff(self.fnorm(query))
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class IsolatedDecisionHead(nn.Module):
|
| 29 |
+
def __init__(self, hidden: int, width: int, heads: int, ff: int, depth: int):
|
| 30 |
+
super().__init__()
|
| 31 |
+
self.norm = nn.LayerNorm(hidden)
|
| 32 |
+
self.memory = nn.Linear(hidden, width, bias=False)
|
| 33 |
+
self.qproj = nn.Linear(hidden, width, bias=False)
|
| 34 |
+
self.oqproj = nn.Linear(hidden, width, bias=False)
|
| 35 |
+
self.globalproj = nn.Linear(hidden, width, bias=False)
|
| 36 |
+
self.context = nn.Linear(hidden, width, bias=False)
|
| 37 |
+
self.lexical = nn.Linear(hidden, width, bias=False)
|
| 38 |
+
self.types = nn.Embedding(3, width)
|
| 39 |
+
self.evidence = nn.ModuleList([EvidenceBlock(width, heads, ff) for _ in range(depth)])
|
| 40 |
+
self.summarynorm = nn.LayerNorm(width)
|
| 41 |
+
self.joint = nn.ModuleList(
|
| 42 |
+
[
|
| 43 |
+
nn.TransformerDecoderLayer(
|
| 44 |
+
width, heads, ff, dropout=0.0, activation="gelu", batch_first=True, norm_first=True
|
| 45 |
+
)
|
| 46 |
+
for _ in range(depth)
|
| 47 |
+
]
|
| 48 |
+
)
|
| 49 |
+
self.fnorm = nn.LayerNorm(width)
|
| 50 |
+
self.onorm = nn.LayerNorm(width)
|
| 51 |
+
self.scorer = nn.Sequential(nn.Linear(width * 4, width), nn.GELU(), nn.Linear(width, 1))
|
| 52 |
+
self.prior_scale = nn.Parameter(torch.zeros(()))
|
| 53 |
+
self.joint_scale = nn.Parameter(torch.zeros(()))
|
| 54 |
+
self.gate = nn.Parameter(torch.zeros(()))
|
| 55 |
+
|
| 56 |
+
def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
|
| 57 |
+
memory = self.memory(self.norm(batch["memory"]))
|
| 58 |
+
question = self.norm(batch["questions"])
|
| 59 |
+
context = self.norm(batch["options"])
|
| 60 |
+
global_feature = self.norm(batch["global"])
|
| 61 |
+
options = self.context(context) + self.lexical(batch["lexical"]) + self.oqproj(question)
|
| 62 |
+
for block in self.evidence:
|
| 63 |
+
options = block(options, memory, batch["memory_pad"])
|
| 64 |
+
fields = self.qproj(question)
|
| 65 |
+
membership = ~batch["option_pad"][:, None, :]
|
| 66 |
+
weights = torch.einsum("bqd,bcd->bqc", fields, options) / memory.shape[-1] ** 0.5
|
| 67 |
+
weights = weights.masked_fill(~membership, -1e9).softmax(-1) * membership
|
| 68 |
+
summary = torch.einsum("bqc,bcd->bqd", weights, options)
|
| 69 |
+
fields = fields + self.summarynorm(summary) + self.globalproj(global_feature)[:, None, :]
|
| 70 |
+
fields = fields + self.types(batch["types"])
|
| 71 |
+
for block in self.joint:
|
| 72 |
+
fields = block(fields, memory, memory_key_padding_mask=batch["memory_pad"])
|
| 73 |
+
fields = self.fnorm(fields).expand(-1, options.shape[1], -1)
|
| 74 |
+
options = self.onorm(options)
|
| 75 |
+
prior = (
|
| 76 |
+
F.normalize(question + global_feature[:, None, :], dim=-1) * F.normalize(batch["lexical"], dim=-1)
|
| 77 |
+
).sum(-1) * self.prior_scale.clamp(max=4.605170186).exp()
|
| 78 |
+
joint = F.cosine_similarity(fields, options, dim=-1) * self.joint_scale.clamp(max=4.605170186).exp()
|
| 79 |
+
residual = self.scorer(torch.cat([fields, options, fields * options, (fields - options).abs()], -1)).squeeze(-1)
|
| 80 |
+
return prior + self.gate.sigmoid() * (joint + residual)
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
@dataclass(frozen=True)
|
| 84 |
+
class QuestionFeatures:
|
| 85 |
+
memory: torch.Tensor
|
| 86 |
+
question: torch.Tensor
|
| 87 |
+
options: torch.Tensor
|
| 88 |
+
lexical: torch.Tensor
|
| 89 |
+
global_feature: torch.Tensor
|
| 90 |
+
question_type: int
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def score_questions(head: IsolatedDecisionHead, records: list[QuestionFeatures]) -> list[torch.Tensor]:
|
| 94 |
+
"""Score all questions in a request batch with one head invocation."""
|
| 95 |
+
if not records:
|
| 96 |
+
raise ValueError("At least one question is required")
|
| 97 |
+
device = head.norm.weight.device
|
| 98 |
+
hidden = head.norm.normalized_shape[0]
|
| 99 |
+
memories = max(len(record.memory) for record in records)
|
| 100 |
+
candidates = max(len(record.options) for record in records)
|
| 101 |
+
batch = {
|
| 102 |
+
name: torch.zeros(shape, device=device, dtype=torch.float32)
|
| 103 |
+
for name, shape in {
|
| 104 |
+
"memory": (len(records), memories, hidden),
|
| 105 |
+
"questions": (len(records), 1, hidden),
|
| 106 |
+
"options": (len(records), candidates, hidden),
|
| 107 |
+
"lexical": (len(records), candidates, hidden),
|
| 108 |
+
"global": (len(records), hidden),
|
| 109 |
+
}.items()
|
| 110 |
+
}
|
| 111 |
+
batch["memory_pad"] = torch.ones(len(records), memories, device=device, dtype=torch.bool)
|
| 112 |
+
batch["option_pad"] = torch.ones(len(records), candidates, device=device, dtype=torch.bool)
|
| 113 |
+
batch["types"] = torch.tensor([[record.question_type] for record in records], device=device)
|
| 114 |
+
for index, record in enumerate(records):
|
| 115 |
+
if not len(record.memory) or not len(record.options) or record.question_type not in (0, 1, 2):
|
| 116 |
+
raise ValueError("Invalid question features")
|
| 117 |
+
if record.lexical.shape != record.options.shape:
|
| 118 |
+
raise ValueError("Criterion embeddings must match candidate features")
|
| 119 |
+
for name, value in (("memory", record.memory), ("options", record.options), ("lexical", record.lexical)):
|
| 120 |
+
batch[name][index, : len(value)] = value.to(device=device, dtype=torch.float32)
|
| 121 |
+
batch["questions"][index, 0] = record.question.to(device=device, dtype=torch.float32)
|
| 122 |
+
batch["global"][index] = record.global_feature.to(device=device, dtype=torch.float32)
|
| 123 |
+
batch["memory_pad"][index, : len(record.memory)] = False
|
| 124 |
+
batch["option_pad"][index, : len(record.options)] = False
|
| 125 |
+
values = head(batch)
|
| 126 |
+
return [values[index, : len(record.options)] for index, record in enumerate(records)]
|
decision_layout.py
ADDED
|
@@ -0,0 +1,84 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Token layout for an evidence prefix and independent decision candidates."""
|
| 2 |
+
|
| 3 |
+
from dataclasses import dataclass
|
| 4 |
+
|
| 5 |
+
from .candidate_tree import CandidateTree
|
| 6 |
+
from .decision_utils import MARKER, literal, locate_markers
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
@dataclass(frozen=True)
|
| 10 |
+
class DecisionLayout:
|
| 11 |
+
tree: CandidateTree
|
| 12 |
+
anchors: tuple[int, ...]
|
| 13 |
+
question_positions: tuple[tuple[int, ...], ...]
|
| 14 |
+
lexical_ids: tuple[tuple[int, ...], ...]
|
| 15 |
+
candidate_questions: tuple[int, ...]
|
| 16 |
+
question_types: tuple[int, ...]
|
| 17 |
+
|
| 18 |
+
def to_dict(self) -> dict:
|
| 19 |
+
return {
|
| 20 |
+
"prefix_length": self.tree.prefix_length,
|
| 21 |
+
"branch_lengths": self.tree.branch_lengths,
|
| 22 |
+
"anchors": self.anchors,
|
| 23 |
+
"question_positions": self.question_positions,
|
| 24 |
+
"lexical_ids": self.lexical_ids,
|
| 25 |
+
"candidate_questions": self.candidate_questions,
|
| 26 |
+
"question_types": self.question_types,
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
@classmethod
|
| 30 |
+
def from_dict(cls, data: dict):
|
| 31 |
+
return cls(
|
| 32 |
+
CandidateTree(data["prefix_length"], tuple(data["branch_lengths"])),
|
| 33 |
+
tuple(data["anchors"]), tuple(tuple(row) for row in data["question_positions"]),
|
| 34 |
+
tuple(tuple(row) for row in data["lexical_ids"]),
|
| 35 |
+
tuple(data["candidate_questions"]), tuple(data["question_types"]),
|
| 36 |
+
)
|
| 37 |
+
|
| 38 |
+
def validate(self, length: int):
|
| 39 |
+
count = len(self.tree.branch_lengths)
|
| 40 |
+
if self.tree.length != length or not (len(self.anchors) == len(self.lexical_ids) == len(self.candidate_questions) == count):
|
| 41 |
+
raise ValueError("Decision layout does not match the packed sequence")
|
| 42 |
+
if not self.question_types or len(self.question_positions) != len(self.question_types):
|
| 43 |
+
raise ValueError("Decision question metadata is incomplete")
|
| 44 |
+
if any(kind not in (0, 1, 2) for kind in self.question_types):
|
| 45 |
+
raise ValueError("Invalid decision question type")
|
| 46 |
+
if set(self.candidate_questions) != set(range(len(self.question_types))):
|
| 47 |
+
raise ValueError("Every question requires candidate branches")
|
| 48 |
+
start = self.tree.prefix_length
|
| 49 |
+
for index, size in enumerate(self.tree.branch_lengths):
|
| 50 |
+
if not start <= self.anchors[index] < start + size or not self.lexical_ids[index]:
|
| 51 |
+
raise ValueError("Invalid decision candidate metadata")
|
| 52 |
+
start += size
|
| 53 |
+
for index, positions in enumerate(self.question_positions):
|
| 54 |
+
branch = self.candidate_questions.index(index)
|
| 55 |
+
start = self.tree.prefix_length + sum(self.tree.branch_lengths[:branch])
|
| 56 |
+
end = start + self.tree.branch_lengths[branch]
|
| 57 |
+
if not positions or any(position < start or position >= end for position in positions):
|
| 58 |
+
raise ValueError("Question positions must belong to its first candidate branch")
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def pack_candidates(tokenizer, prefix: list[int], questions: dict) -> tuple[list[int], DecisionLayout]:
|
| 62 |
+
ids = list(prefix)
|
| 63 |
+
lengths, anchors, positions, lexical, owners, kinds = [], [], [], [], [], []
|
| 64 |
+
marker = tokenizer.encode(MARKER.lstrip("\n"), add_special_tokens=False)
|
| 65 |
+
for question_index, question in enumerate(questions.values()):
|
| 66 |
+
kinds.append({"predicate": 0, "choice": 1, "score": 2}[question["type"]])
|
| 67 |
+
for candidate_index, description in enumerate(question["criteria"].values()):
|
| 68 |
+
text = literal({"question": question["instructions"], "criterion": description}) + MARKER
|
| 69 |
+
encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
|
| 70 |
+
start = len(ids)
|
| 71 |
+
branch = encoded["input_ids"]
|
| 72 |
+
ids.extend(branch)
|
| 73 |
+
lengths.append(len(branch))
|
| 74 |
+
anchors.append(start + locate_markers(branch, marker, 1)[0])
|
| 75 |
+
owners.append(question_index)
|
| 76 |
+
lexical.append(tuple(tokenizer.encode(literal(description), add_special_tokens=False)))
|
| 77 |
+
if candidate_index == 0:
|
| 78 |
+
phrase = literal(question["instructions"])
|
| 79 |
+
begin = len('{"question": ')
|
| 80 |
+
end = begin + len(phrase)
|
| 81 |
+
positions.append(tuple(start + i for i, (a, b) in enumerate(encoded["offset_mapping"]) if a < end and b > begin))
|
| 82 |
+
layout = DecisionLayout(CandidateTree(len(prefix), tuple(lengths)), tuple(anchors), tuple(positions), tuple(lexical), tuple(owners), tuple(kinds))
|
| 83 |
+
layout.validate(len(ids))
|
| 84 |
+
return ids, layout
|
decision_utils.py
CHANGED
|
@@ -13,7 +13,7 @@ MARKER = "\n<spev_score>\n"
|
|
| 13 |
|
| 14 |
class ChoiceQuestion(BaseModel):
|
| 15 |
model_config = ConfigDict(extra="forbid", strict=True)
|
| 16 |
-
type: Literal["choice"] = "choice"
|
| 17 |
instructions: str
|
| 18 |
criteria: dict[str, str] = Field(min_length=2, max_length=255)
|
| 19 |
|
|
|
|
| 13 |
|
| 14 |
class ChoiceQuestion(BaseModel):
|
| 15 |
model_config = ConfigDict(extra="forbid", strict=True)
|
| 16 |
+
type: Literal["choice", "predicate", "score"] = "choice"
|
| 17 |
instructions: str
|
| 18 |
criteria: dict[str, str] = Field(min_length=2, max_length=255)
|
| 19 |
|
examples/turn_decision.py
DELETED
|
@@ -1,28 +0,0 @@
|
|
| 1 |
-
"""Evaluate an audio file using the packaged turn-taking question."""
|
| 2 |
-
|
| 3 |
-
import argparse
|
| 4 |
-
import json
|
| 5 |
-
from pathlib import Path
|
| 6 |
-
|
| 7 |
-
import torch
|
| 8 |
-
from transformers import AutoModel, AutoProcessor
|
| 9 |
-
|
| 10 |
-
parser = argparse.ArgumentParser()
|
| 11 |
-
parser.add_argument("audio")
|
| 12 |
-
parser.add_argument("--model", default=str(Path(__file__).resolve().parents[1]))
|
| 13 |
-
args = parser.parse_args()
|
| 14 |
-
policy = json.loads((Path(__file__).with_name("turn_policy.json")).read_text())
|
| 15 |
-
processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
|
| 16 |
-
model = AutoModel.from_pretrained(
|
| 17 |
-
args.model,
|
| 18 |
-
trust_remote_code=True,
|
| 19 |
-
dtype=torch.bfloat16,
|
| 20 |
-
device_map="auto",
|
| 21 |
-
attn_implementation="sdpa",
|
| 22 |
-
)
|
| 23 |
-
inputs = processor(audio=args.audio, state={}, questions=policy["questions"])
|
| 24 |
-
inputs = inputs.to(model.device, dtype=torch.bfloat16)
|
| 25 |
-
with torch.inference_mode():
|
| 26 |
-
outputs = model(**inputs)
|
| 27 |
-
result = processor.post_process(outputs, questions=policy["questions"])
|
| 28 |
-
print(json.dumps(result, indent=2))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/turn_policy.json
DELETED
|
@@ -1,17 +0,0 @@
|
|
| 1 |
-
{
|
| 2 |
-
"questions": {
|
| 3 |
-
"decision": {
|
| 4 |
-
"instructions": "At the end of this conversation audio, has the most recent speaker yielded the floor or are they continuing their turn?",
|
| 5 |
-
"criteria": {
|
| 6 |
-
"keep_listening": "The most recent speaker is continuing their turn; keep listening.",
|
| 7 |
-
"respond": "The most recent speaker has yielded the floor; another speaker can respond."
|
| 8 |
-
}
|
| 9 |
-
}
|
| 10 |
-
},
|
| 11 |
-
"temperature": 1.3234945284875692,
|
| 12 |
-
"response_threshold": 0.3856703463365996,
|
| 13 |
-
"comparison": "strictly greater than",
|
| 14 |
-
"scope": "Fixed English turn-taking schema; not a generic Choice threshold",
|
| 15 |
-
"temperature_fitted_on": "AMI designated calibration split",
|
| 16 |
-
"threshold_selected_on": "AMI development split"
|
| 17 |
-
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
media.py
ADDED
|
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Bounded video sampling for SPEV evidence."""
|
| 2 |
+
|
| 3 |
+
import io
|
| 4 |
+
import math
|
| 5 |
+
from tempfile import NamedTemporaryFile
|
| 6 |
+
|
| 7 |
+
import av
|
| 8 |
+
from transformers.video_utils import load_video
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def decode_video(data, video_processor, *, fps, max_seconds, max_frames, max_bytes):
|
| 12 |
+
if len(data) > max_bytes:
|
| 13 |
+
raise ValueError("Video exceeds the checkpoint byte limit")
|
| 14 |
+
try:
|
| 15 |
+
with av.open(io.BytesIO(data)) as container:
|
| 16 |
+
if not container.streams.video:
|
| 17 |
+
raise ValueError("Video has no visual stream")
|
| 18 |
+
stream = container.streams.video[0]
|
| 19 |
+
duration = float(stream.duration * stream.time_base) if stream.duration is not None else None
|
| 20 |
+
if duration is None and container.duration is not None:
|
| 21 |
+
duration = container.duration / av.time_base
|
| 22 |
+
if duration is None or not math.isfinite(duration) or not 0 < duration <= max_seconds:
|
| 23 |
+
raise ValueError("Video duration is missing or exceeds the checkpoint limit")
|
| 24 |
+
frame_rate = float(stream.average_rate) if stream.average_rate is not None else None
|
| 25 |
+
if frame_rate is None or not math.isfinite(frame_rate) or frame_rate <= 0:
|
| 26 |
+
raise ValueError("Video frame rate is missing or invalid")
|
| 27 |
+
frame_count = stream.frames or math.ceil(duration * frame_rate)
|
| 28 |
+
if not 0 < frame_count <= max_frames:
|
| 29 |
+
raise ValueError("Video frame count is missing or exceeds the checkpoint limit")
|
| 30 |
+
except av.FFmpegError as exc:
|
| 31 |
+
raise ValueError("Invalid or unsupported video") from exc
|
| 32 |
+
|
| 33 |
+
def sample_indices(metadata, **kwargs):
|
| 34 |
+
metadata.total_num_frames = frame_count
|
| 35 |
+
metadata.fps = frame_rate
|
| 36 |
+
metadata.duration = frame_count / frame_rate
|
| 37 |
+
return video_processor.sample_frames(metadata=metadata, fps=fps)
|
| 38 |
+
|
| 39 |
+
with NamedTemporaryFile() as source:
|
| 40 |
+
source.write(data)
|
| 41 |
+
source.flush()
|
| 42 |
+
try:
|
| 43 |
+
frames, metadata = load_video(source.name, backend="pyav", sample_indices_fn=sample_indices)
|
| 44 |
+
except (av.FFmpegError, OSError) as exc:
|
| 45 |
+
raise ValueError("Invalid or unsupported video") from exc
|
| 46 |
+
if not len(frames) or not metadata.fps or not metadata.total_num_frames:
|
| 47 |
+
raise ValueError("Video has no decodable frames")
|
| 48 |
+
if metadata.frames_indices is None or len(frames) != len(metadata.frames_indices):
|
| 49 |
+
raise ValueError("Video sampled-frame decode is incomplete")
|
| 50 |
+
return frames, metadata
|
model-decision.safetensors
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ea005b67b23edcb83fe48a98978452dfd382db9937c789ba445236a6e9feb6a2
|
| 3 |
+
size 81916464
|
model.safetensors.index.json
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
{
|
| 2 |
"metadata": {
|
| 3 |
-
"total_size":
|
| 4 |
},
|
| 5 |
"weight_map": {
|
| 6 |
"backbone.audio_tower.conv2d1.bias": "model-00001-of-00013.safetensors",
|
|
@@ -19937,11 +19937,91 @@
|
|
| 19937 |
"backbone.model.layers.9.self_attn.q_proj.b": "model-decision.safetensors",
|
| 19938 |
"backbone.model.layers.9.self_attn.v_proj.a": "model-decision.safetensors",
|
| 19939 |
"backbone.model.layers.9.self_attn.v_proj.b": "model-decision.safetensors",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19940 |
"head.scorer.0.bias": "model-decision.safetensors",
|
| 19941 |
"head.scorer.0.weight": "model-decision.safetensors",
|
| 19942 |
-
"head.scorer.
|
| 19943 |
-
"head.scorer.
|
| 19944 |
-
"head.
|
| 19945 |
-
"head.
|
|
|
|
| 19946 |
}
|
| 19947 |
}
|
|
|
|
| 1 |
{
|
| 2 |
"metadata": {
|
| 3 |
+
"total_size": 62897443056
|
| 4 |
},
|
| 5 |
"weight_map": {
|
| 6 |
"backbone.audio_tower.conv2d1.bias": "model-00001-of-00013.safetensors",
|
|
|
|
| 19937 |
"backbone.model.layers.9.self_attn.q_proj.b": "model-decision.safetensors",
|
| 19938 |
"backbone.model.layers.9.self_attn.v_proj.a": "model-decision.safetensors",
|
| 19939 |
"backbone.model.layers.9.self_attn.v_proj.b": "model-decision.safetensors",
|
| 19940 |
+
"head.context.weight": "model-decision.safetensors",
|
| 19941 |
+
"head.evidence.0.attn.in_proj_bias": "model-decision.safetensors",
|
| 19942 |
+
"head.evidence.0.attn.in_proj_weight": "model-decision.safetensors",
|
| 19943 |
+
"head.evidence.0.attn.out_proj.bias": "model-decision.safetensors",
|
| 19944 |
+
"head.evidence.0.attn.out_proj.weight": "model-decision.safetensors",
|
| 19945 |
+
"head.evidence.0.ff.0.bias": "model-decision.safetensors",
|
| 19946 |
+
"head.evidence.0.ff.0.weight": "model-decision.safetensors",
|
| 19947 |
+
"head.evidence.0.ff.2.bias": "model-decision.safetensors",
|
| 19948 |
+
"head.evidence.0.ff.2.weight": "model-decision.safetensors",
|
| 19949 |
+
"head.evidence.0.fnorm.bias": "model-decision.safetensors",
|
| 19950 |
+
"head.evidence.0.fnorm.weight": "model-decision.safetensors",
|
| 19951 |
+
"head.evidence.0.mnorm.bias": "model-decision.safetensors",
|
| 19952 |
+
"head.evidence.0.mnorm.weight": "model-decision.safetensors",
|
| 19953 |
+
"head.evidence.0.qnorm.bias": "model-decision.safetensors",
|
| 19954 |
+
"head.evidence.0.qnorm.weight": "model-decision.safetensors",
|
| 19955 |
+
"head.evidence.1.attn.in_proj_bias": "model-decision.safetensors",
|
| 19956 |
+
"head.evidence.1.attn.in_proj_weight": "model-decision.safetensors",
|
| 19957 |
+
"head.evidence.1.attn.out_proj.bias": "model-decision.safetensors",
|
| 19958 |
+
"head.evidence.1.attn.out_proj.weight": "model-decision.safetensors",
|
| 19959 |
+
"head.evidence.1.ff.0.bias": "model-decision.safetensors",
|
| 19960 |
+
"head.evidence.1.ff.0.weight": "model-decision.safetensors",
|
| 19961 |
+
"head.evidence.1.ff.2.bias": "model-decision.safetensors",
|
| 19962 |
+
"head.evidence.1.ff.2.weight": "model-decision.safetensors",
|
| 19963 |
+
"head.evidence.1.fnorm.bias": "model-decision.safetensors",
|
| 19964 |
+
"head.evidence.1.fnorm.weight": "model-decision.safetensors",
|
| 19965 |
+
"head.evidence.1.mnorm.bias": "model-decision.safetensors",
|
| 19966 |
+
"head.evidence.1.mnorm.weight": "model-decision.safetensors",
|
| 19967 |
+
"head.evidence.1.qnorm.bias": "model-decision.safetensors",
|
| 19968 |
+
"head.evidence.1.qnorm.weight": "model-decision.safetensors",
|
| 19969 |
+
"head.fnorm.bias": "model-decision.safetensors",
|
| 19970 |
+
"head.fnorm.weight": "model-decision.safetensors",
|
| 19971 |
+
"head.gate": "model-decision.safetensors",
|
| 19972 |
+
"head.globalproj.weight": "model-decision.safetensors",
|
| 19973 |
+
"head.joint.0.linear1.bias": "model-decision.safetensors",
|
| 19974 |
+
"head.joint.0.linear1.weight": "model-decision.safetensors",
|
| 19975 |
+
"head.joint.0.linear2.bias": "model-decision.safetensors",
|
| 19976 |
+
"head.joint.0.linear2.weight": "model-decision.safetensors",
|
| 19977 |
+
"head.joint.0.multihead_attn.in_proj_bias": "model-decision.safetensors",
|
| 19978 |
+
"head.joint.0.multihead_attn.in_proj_weight": "model-decision.safetensors",
|
| 19979 |
+
"head.joint.0.multihead_attn.out_proj.bias": "model-decision.safetensors",
|
| 19980 |
+
"head.joint.0.multihead_attn.out_proj.weight": "model-decision.safetensors",
|
| 19981 |
+
"head.joint.0.norm1.bias": "model-decision.safetensors",
|
| 19982 |
+
"head.joint.0.norm1.weight": "model-decision.safetensors",
|
| 19983 |
+
"head.joint.0.norm2.bias": "model-decision.safetensors",
|
| 19984 |
+
"head.joint.0.norm2.weight": "model-decision.safetensors",
|
| 19985 |
+
"head.joint.0.norm3.bias": "model-decision.safetensors",
|
| 19986 |
+
"head.joint.0.norm3.weight": "model-decision.safetensors",
|
| 19987 |
+
"head.joint.0.self_attn.in_proj_bias": "model-decision.safetensors",
|
| 19988 |
+
"head.joint.0.self_attn.in_proj_weight": "model-decision.safetensors",
|
| 19989 |
+
"head.joint.0.self_attn.out_proj.bias": "model-decision.safetensors",
|
| 19990 |
+
"head.joint.0.self_attn.out_proj.weight": "model-decision.safetensors",
|
| 19991 |
+
"head.joint.1.linear1.bias": "model-decision.safetensors",
|
| 19992 |
+
"head.joint.1.linear1.weight": "model-decision.safetensors",
|
| 19993 |
+
"head.joint.1.linear2.bias": "model-decision.safetensors",
|
| 19994 |
+
"head.joint.1.linear2.weight": "model-decision.safetensors",
|
| 19995 |
+
"head.joint.1.multihead_attn.in_proj_bias": "model-decision.safetensors",
|
| 19996 |
+
"head.joint.1.multihead_attn.in_proj_weight": "model-decision.safetensors",
|
| 19997 |
+
"head.joint.1.multihead_attn.out_proj.bias": "model-decision.safetensors",
|
| 19998 |
+
"head.joint.1.multihead_attn.out_proj.weight": "model-decision.safetensors",
|
| 19999 |
+
"head.joint.1.norm1.bias": "model-decision.safetensors",
|
| 20000 |
+
"head.joint.1.norm1.weight": "model-decision.safetensors",
|
| 20001 |
+
"head.joint.1.norm2.bias": "model-decision.safetensors",
|
| 20002 |
+
"head.joint.1.norm2.weight": "model-decision.safetensors",
|
| 20003 |
+
"head.joint.1.norm3.bias": "model-decision.safetensors",
|
| 20004 |
+
"head.joint.1.norm3.weight": "model-decision.safetensors",
|
| 20005 |
+
"head.joint.1.self_attn.in_proj_bias": "model-decision.safetensors",
|
| 20006 |
+
"head.joint.1.self_attn.in_proj_weight": "model-decision.safetensors",
|
| 20007 |
+
"head.joint.1.self_attn.out_proj.bias": "model-decision.safetensors",
|
| 20008 |
+
"head.joint.1.self_attn.out_proj.weight": "model-decision.safetensors",
|
| 20009 |
+
"head.joint_scale": "model-decision.safetensors",
|
| 20010 |
+
"head.lexical.weight": "model-decision.safetensors",
|
| 20011 |
+
"head.memory.weight": "model-decision.safetensors",
|
| 20012 |
+
"head.norm.bias": "model-decision.safetensors",
|
| 20013 |
+
"head.norm.weight": "model-decision.safetensors",
|
| 20014 |
+
"head.onorm.bias": "model-decision.safetensors",
|
| 20015 |
+
"head.onorm.weight": "model-decision.safetensors",
|
| 20016 |
+
"head.oqproj.weight": "model-decision.safetensors",
|
| 20017 |
+
"head.prior_scale": "model-decision.safetensors",
|
| 20018 |
+
"head.qproj.weight": "model-decision.safetensors",
|
| 20019 |
"head.scorer.0.bias": "model-decision.safetensors",
|
| 20020 |
"head.scorer.0.weight": "model-decision.safetensors",
|
| 20021 |
+
"head.scorer.2.bias": "model-decision.safetensors",
|
| 20022 |
+
"head.scorer.2.weight": "model-decision.safetensors",
|
| 20023 |
+
"head.summarynorm.bias": "model-decision.safetensors",
|
| 20024 |
+
"head.summarynorm.weight": "model-decision.safetensors",
|
| 20025 |
+
"head.types.weight": "model-decision.safetensors"
|
| 20026 |
}
|
| 20027 |
}
|
modeling_spev.py
CHANGED
|
@@ -10,6 +10,10 @@ from transformers import PreTrainedModel, Qwen3OmniMoeThinkerForConditionalGener
|
|
| 10 |
from transformers.utils import ModelOutput
|
| 11 |
|
| 12 |
from .configuration_spev import SPEVConfig
|
|
|
|
|
|
|
|
|
|
|
|
|
| 13 |
|
| 14 |
|
| 15 |
class ResidualAdapter(nn.Module):
|
|
@@ -26,20 +30,6 @@ class ResidualAdapter(nn.Module):
|
|
| 26 |
return self.base(x) + (residual * self.scale).to(x.dtype)
|
| 27 |
|
| 28 |
|
| 29 |
-
class OptionHead(nn.Module):
|
| 30 |
-
def __init__(self, hidden_size, width):
|
| 31 |
-
super().__init__()
|
| 32 |
-
self.scorer = nn.Sequential(
|
| 33 |
-
nn.LayerNorm(hidden_size),
|
| 34 |
-
nn.Linear(hidden_size, width),
|
| 35 |
-
nn.GELU(),
|
| 36 |
-
nn.Linear(width, 1),
|
| 37 |
-
)
|
| 38 |
-
|
| 39 |
-
def forward(self, features):
|
| 40 |
-
return self.scorer(features.float()).squeeze(-1)
|
| 41 |
-
|
| 42 |
-
|
| 43 |
@dataclass
|
| 44 |
class SPEVOutput(ModelOutput):
|
| 45 |
logits: torch.Tensor | None = None
|
|
@@ -57,7 +47,7 @@ class SPEVModel(PreTrainedModel):
|
|
| 57 |
_no_split_modules: ClassVar[list[str]] = [
|
| 58 |
"Qwen3OmniMoeThinkerTextDecoderLayer",
|
| 59 |
"ResidualAdapter",
|
| 60 |
-
"
|
| 61 |
]
|
| 62 |
|
| 63 |
def __init__(self, config):
|
|
@@ -76,10 +66,13 @@ class SPEVModel(PreTrainedModel):
|
|
| 76 |
leaf,
|
| 77 |
ResidualAdapter(module, config.adapter_rank, config.adapter_alpha),
|
| 78 |
)
|
| 79 |
-
self.head =
|
| 80 |
-
config.thinker_config.text_config.hidden_size, config.head_width
|
|
|
|
| 81 |
)
|
| 82 |
self.post_init()
|
|
|
|
|
|
|
| 83 |
|
| 84 |
def get_input_embeddings(self):
|
| 85 |
return self.backbone.get_input_embeddings()
|
|
@@ -88,46 +81,70 @@ class SPEVModel(PreTrainedModel):
|
|
| 88 |
self.backbone.set_input_embeddings(value)
|
| 89 |
|
| 90 |
def forward(
|
| 91 |
-
self, input_ids=None, attention_mask=None,
|
| 92 |
):
|
| 93 |
if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] != 1:
|
| 94 |
raise ValueError("SPEV accepts one evidence sequence per forward pass")
|
| 95 |
if input_ids.shape[1] > self.config.max_tokens:
|
| 96 |
raise ValueError("Input exceeds the configured context limit")
|
| 97 |
-
if (
|
| 98 |
-
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
|
| 103 |
-
|
| 104 |
-
scoring_positions.numel() == 0
|
| 105 |
-
or torch.any(scoring_positions < 0)
|
| 106 |
-
or torch.any(scoring_positions >= input_ids.shape[1])
|
| 107 |
-
):
|
| 108 |
-
raise ValueError("Scoring positions are outside the input sequence")
|
| 109 |
kwargs.pop("use_cache", None)
|
| 110 |
kwargs.pop("return_dict", None)
|
| 111 |
-
|
| 112 |
-
|
| 113 |
-
|
| 114 |
-
|
| 115 |
-
|
| 116 |
-
|
| 117 |
-
|
| 118 |
-
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
head_device = self.head.scorer[0].weight.device
|
| 123 |
-
logits = self.head(features.to(device=head_device, dtype=torch.float32))
|
| 124 |
-
count = (
|
| 125 |
-
attention_mask.sum(-1)
|
| 126 |
-
if attention_mask is not None
|
| 127 |
-
else input_ids.new_tensor([input_ids.shape[1]])
|
| 128 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 129 |
return SPEVOutput(
|
| 130 |
-
logits=logits, input_tokens=
|
|
|
|
| 131 |
)
|
| 132 |
|
| 133 |
|
|
|
|
| 10 |
from transformers.utils import ModelOutput
|
| 11 |
|
| 12 |
from .configuration_spev import SPEVConfig
|
| 13 |
+
from .branch_attention import install as install_attention, request_context
|
| 14 |
+
from .decision_head import IsolatedDecisionHead, QuestionFeatures, score_questions
|
| 15 |
+
from .decision_layout import DecisionLayout
|
| 16 |
+
from .token_operations import install as install_token_operations
|
| 17 |
|
| 18 |
|
| 19 |
class ResidualAdapter(nn.Module):
|
|
|
|
| 30 |
return self.base(x) + (residual * self.scale).to(x.dtype)
|
| 31 |
|
| 32 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
@dataclass
|
| 34 |
class SPEVOutput(ModelOutput):
|
| 35 |
logits: torch.Tensor | None = None
|
|
|
|
| 47 |
_no_split_modules: ClassVar[list[str]] = [
|
| 48 |
"Qwen3OmniMoeThinkerTextDecoderLayer",
|
| 49 |
"ResidualAdapter",
|
| 50 |
+
"IsolatedDecisionHead",
|
| 51 |
]
|
| 52 |
|
| 53 |
def __init__(self, config):
|
|
|
|
| 66 |
leaf,
|
| 67 |
ResidualAdapter(module, config.adapter_rank, config.adapter_alpha),
|
| 68 |
)
|
| 69 |
+
self.head = IsolatedDecisionHead(
|
| 70 |
+
config.thinker_config.text_config.hidden_size, config.head_width,
|
| 71 |
+
config.head_heads, config.head_ff, config.head_depth,
|
| 72 |
)
|
| 73 |
self.post_init()
|
| 74 |
+
install_attention(self.backbone)
|
| 75 |
+
install_token_operations(self, linears=True, moe=False)
|
| 76 |
|
| 77 |
def get_input_embeddings(self):
|
| 78 |
return self.backbone.get_input_embeddings()
|
|
|
|
| 81 |
self.backbone.set_input_embeddings(value)
|
| 82 |
|
| 83 |
def forward(
|
| 84 |
+
self, input_ids=None, attention_mask=None, decision_layout=None, **kwargs
|
| 85 |
):
|
| 86 |
if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] != 1:
|
| 87 |
raise ValueError("SPEV accepts one evidence sequence per forward pass")
|
| 88 |
if input_ids.shape[1] > self.config.max_tokens:
|
| 89 |
raise ValueError("Input exceeds the configured context limit")
|
| 90 |
+
if not isinstance(decision_layout, dict):
|
| 91 |
+
raise ValueError("A processor-generated decision_layout is required")
|
| 92 |
+
decision_layout = DecisionLayout.from_dict(decision_layout)
|
| 93 |
+
decision_layout.validate(input_ids.shape[1])
|
| 94 |
+
tree = decision_layout.tree
|
| 95 |
+
if attention_mask is not None and not torch.all(attention_mask == 1):
|
| 96 |
+
raise ValueError("SPEV accepts unpadded evidence sequences")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 97 |
kwargs.pop("use_cache", None)
|
| 98 |
kwargs.pop("return_dict", None)
|
| 99 |
+
kwargs.pop("position_ids", None)
|
| 100 |
+
prefix = input_ids[:, :tree.prefix_length]
|
| 101 |
+
rope_kwargs = {key: kwargs[key] for key in (
|
| 102 |
+
"image_grid_thw", "video_grid_thw", "use_audio_in_video", "second_per_grids"
|
| 103 |
+
) if key in kwargs}
|
| 104 |
+
if "video_second_per_grid" in kwargs:
|
| 105 |
+
rope_kwargs["second_per_grids"] = kwargs["video_second_per_grid"]
|
| 106 |
+
if "feature_attention_mask" in kwargs:
|
| 107 |
+
rope_kwargs["audio_seqlens"] = kwargs["feature_attention_mask"].sum(1)
|
| 108 |
+
positions, _ = self.backbone.get_rope_index(
|
| 109 |
+
prefix, attention_mask=torch.ones_like(prefix), **rope_kwargs
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 110 |
)
|
| 111 |
+
position_ids = tree.position_ids(positions).to(input_ids.device)
|
| 112 |
+
with request_context(tree):
|
| 113 |
+
hidden = self.backbone(
|
| 114 |
+
input_ids=input_ids,
|
| 115 |
+
attention_mask=tree.additive_mask(dtype=self.backbone.dtype, device=input_ids.device),
|
| 116 |
+
position_ids=position_ids,
|
| 117 |
+
use_cache=False,
|
| 118 |
+
return_dict=True,
|
| 119 |
+
**kwargs,
|
| 120 |
+
).logits[0]
|
| 121 |
+
lexical = []
|
| 122 |
+
for tokens in decision_layout.lexical_ids:
|
| 123 |
+
tokens = input_ids.new_tensor(tokens)
|
| 124 |
+
lexical.append(self.get_input_embeddings()(tokens).mean(0))
|
| 125 |
+
records = []
|
| 126 |
+
starts = [tree.prefix_length]
|
| 127 |
+
for length in tree.branch_lengths:
|
| 128 |
+
starts.append(starts[-1] + length)
|
| 129 |
+
for question, kind in enumerate(decision_layout.question_types):
|
| 130 |
+
indices = [i for i, owner in enumerate(decision_layout.candidate_questions) if owner == question]
|
| 131 |
+
memory = torch.cat([hidden[:tree.prefix_length]] + [hidden[starts[i]:starts[i + 1]] for i in indices])
|
| 132 |
+
records.append(QuestionFeatures(
|
| 133 |
+
memory=memory,
|
| 134 |
+
question=hidden[list(decision_layout.question_positions[question])].mean(0),
|
| 135 |
+
options=hidden[[decision_layout.anchors[i] for i in indices]],
|
| 136 |
+
lexical=torch.stack([lexical[i] for i in indices]),
|
| 137 |
+
global_feature=hidden[tree.prefix_length - 1],
|
| 138 |
+
question_type=kind,
|
| 139 |
+
))
|
| 140 |
+
grouped = score_questions(self.head, records)
|
| 141 |
+
logits = torch.empty(len(lexical), device=grouped[0].device, dtype=grouped[0].dtype)
|
| 142 |
+
for question, values in enumerate(grouped):
|
| 143 |
+
indices = [i for i, owner in enumerate(decision_layout.candidate_questions) if owner == question]
|
| 144 |
+
logits[indices] = values
|
| 145 |
return SPEVOutput(
|
| 146 |
+
logits=logits[None], input_tokens=input_ids.new_tensor([input_ids.shape[1]]),
|
| 147 |
+
temperature=self.config.temperature
|
| 148 |
)
|
| 149 |
|
| 150 |
|
processing_spev.py
CHANGED
|
@@ -7,10 +7,10 @@ from .decision_utils import (
|
|
| 7 |
MARKER,
|
| 8 |
ChoiceQuestions,
|
| 9 |
answers,
|
| 10 |
-
locate_markers,
|
| 11 |
prepare_audio,
|
| 12 |
-
|
| 13 |
)
|
|
|
|
| 14 |
|
| 15 |
|
| 16 |
class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
@@ -23,6 +23,9 @@ class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
| 23 |
chat_template=None,
|
| 24 |
max_tokens=8192,
|
| 25 |
video_fps=1.0,
|
|
|
|
|
|
|
|
|
|
| 26 |
):
|
| 27 |
super().__init__(
|
| 28 |
image_processor,
|
|
@@ -33,6 +36,9 @@ class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
| 33 |
)
|
| 34 |
self.max_tokens = max_tokens
|
| 35 |
self.video_fps = video_fps
|
|
|
|
|
|
|
|
|
|
| 36 |
self.marker = self.tokenizer.encode(
|
| 37 |
MARKER.lstrip("\n"), add_special_tokens=False
|
| 38 |
)
|
|
@@ -41,7 +47,8 @@ class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
| 41 |
questions = ChoiceQuestions.model_validate(
|
| 42 |
{"questions": questions}
|
| 43 |
).model_dump()["questions"]
|
| 44 |
-
text
|
|
|
|
| 45 |
content, media = [], {}
|
| 46 |
if audio is not None:
|
| 47 |
import soundfile as sf
|
|
@@ -56,20 +63,17 @@ class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
| 56 |
media["images"] = [source.convert("RGB")]
|
| 57 |
content.append({"type": "image", "image": "provided-image"})
|
| 58 |
if video is not None:
|
| 59 |
-
from
|
|
|
|
| 60 |
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
fps=self.video_fps,
|
| 69 |
-
),
|
| 70 |
)
|
| 71 |
-
if not len(frames) or not metadata.fps or not metadata.total_num_frames:
|
| 72 |
-
raise ValueError("Video has no decodable frames or timing information")
|
| 73 |
media["videos"] = [frames]
|
| 74 |
media["fps"] = len(frames) * metadata.fps / metadata.total_num_frames
|
| 75 |
media["do_sample_frames"] = False
|
|
@@ -93,10 +97,15 @@ class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
| 93 |
**media,
|
| 94 |
)
|
| 95 |
ids = inputs["input_ids"][0].tolist()
|
|
|
|
|
|
|
|
|
|
|
|
|
| 96 |
if len(ids) > self.max_tokens:
|
| 97 |
raise ValueError("Input exceeds the configured context limit")
|
| 98 |
-
|
| 99 |
-
inputs["
|
|
|
|
| 100 |
return inputs
|
| 101 |
|
| 102 |
def post_process(self, outputs, *, questions):
|
|
|
|
| 7 |
MARKER,
|
| 8 |
ChoiceQuestions,
|
| 9 |
answers,
|
|
|
|
| 10 |
prepare_audio,
|
| 11 |
+
literal,
|
| 12 |
)
|
| 13 |
+
from .decision_layout import pack_candidates
|
| 14 |
|
| 15 |
|
| 16 |
class SPEVProcessor(Qwen3OmniMoeProcessor):
|
|
|
|
| 23 |
chat_template=None,
|
| 24 |
max_tokens=8192,
|
| 25 |
video_fps=1.0,
|
| 26 |
+
max_video_seconds=60.0,
|
| 27 |
+
max_video_frames=1800,
|
| 28 |
+
max_media_bytes=12 * 1024 * 1024,
|
| 29 |
):
|
| 30 |
super().__init__(
|
| 31 |
image_processor,
|
|
|
|
| 36 |
)
|
| 37 |
self.max_tokens = max_tokens
|
| 38 |
self.video_fps = video_fps
|
| 39 |
+
self.max_video_seconds = max_video_seconds
|
| 40 |
+
self.max_video_frames = max_video_frames
|
| 41 |
+
self.max_media_bytes = max_media_bytes
|
| 42 |
self.marker = self.tokenizer.encode(
|
| 43 |
MARKER.lstrip("\n"), add_special_tokens=False
|
| 44 |
)
|
|
|
|
| 47 |
questions = ChoiceQuestions.model_validate(
|
| 48 |
{"questions": questions}
|
| 49 |
).model_dump()["questions"]
|
| 50 |
+
text = "Use the audio and state as evidence for the supplied questions.\n"
|
| 51 |
+
text += "State: " + literal({} if state is None else state) + "\n"
|
| 52 |
content, media = [], {}
|
| 53 |
if audio is not None:
|
| 54 |
import soundfile as sf
|
|
|
|
| 63 |
media["images"] = [source.convert("RGB")]
|
| 64 |
content.append({"type": "image", "image": "provided-image"})
|
| 65 |
if video is not None:
|
| 66 |
+
from pathlib import Path
|
| 67 |
+
from .media import decode_video
|
| 68 |
|
| 69 |
+
source = Path(video)
|
| 70 |
+
if source.stat().st_size > self.max_media_bytes:
|
| 71 |
+
raise ValueError("Video exceeds the checkpoint byte limit")
|
| 72 |
+
frames, metadata = decode_video(
|
| 73 |
+
source.read_bytes(), self.video_processor, fps=self.video_fps,
|
| 74 |
+
max_seconds=self.max_video_seconds, max_frames=self.max_video_frames,
|
| 75 |
+
max_bytes=self.max_media_bytes,
|
|
|
|
|
|
|
| 76 |
)
|
|
|
|
|
|
|
| 77 |
media["videos"] = [frames]
|
| 78 |
media["fps"] = len(frames) * metadata.fps / metadata.total_num_frames
|
| 79 |
media["do_sample_frames"] = False
|
|
|
|
| 97 |
**media,
|
| 98 |
)
|
| 99 |
ids = inputs["input_ids"][0].tolist()
|
| 100 |
+
ending = self.tokenizer.encode("<|im_end|>\n", add_special_tokens=False)
|
| 101 |
+
if ids[-len(ending):] != ending:
|
| 102 |
+
raise ValueError("Evidence chat template has an unexpected ending")
|
| 103 |
+
ids, layout = pack_candidates(self.tokenizer, ids[:-len(ending)], questions)
|
| 104 |
if len(ids) > self.max_tokens:
|
| 105 |
raise ValueError("Input exceeds the configured context limit")
|
| 106 |
+
inputs["input_ids"] = torch.tensor([ids], dtype=torch.long)
|
| 107 |
+
inputs["attention_mask"] = torch.ones(1, len(ids), dtype=torch.long)
|
| 108 |
+
inputs["decision_layout"] = layout.to_dict()
|
| 109 |
return inputs
|
| 110 |
|
| 111 |
def post_process(self, outputs, *, questions):
|
processor_config.json
CHANGED
|
@@ -74,7 +74,10 @@
|
|
| 74 |
},
|
| 75 |
"temporal_patch_size": 2
|
| 76 |
},
|
|
|
|
| 77 |
"max_tokens": 8192,
|
|
|
|
|
|
|
| 78 |
"processor_class": "SPEVProcessor",
|
| 79 |
"video_fps": 1.0,
|
| 80 |
"video_processor": {
|
|
|
|
| 74 |
},
|
| 75 |
"temporal_patch_size": 2
|
| 76 |
},
|
| 77 |
+
"max_media_bytes": 12582912,
|
| 78 |
"max_tokens": 8192,
|
| 79 |
+
"max_video_frames": 1800,
|
| 80 |
+
"max_video_seconds": 60.0,
|
| 81 |
"processor_class": "SPEVProcessor",
|
| 82 |
"video_fps": 1.0,
|
| 83 |
"video_processor": {
|
token_operations.py
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Reuse request-local gather indices for fixed-shape token-local operations."""
|
| 2 |
+
from collections import defaultdict
|
| 3 |
+
from .branch_attention import REQUEST
|
| 4 |
+
from contextvars import ContextVar
|
| 5 |
+
import types
|
| 6 |
+
import torch
|
| 7 |
+
|
| 8 |
+
DEPTH=ContextVar('spev_vectorized_token_ops_depth',default=0)
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
class PackingCache:
|
| 12 |
+
def __init__(self):
|
| 13 |
+
self.tree=None;self.device=None;self.plan=None
|
| 14 |
+
|
| 15 |
+
def __call__(self,tree,device):
|
| 16 |
+
if tree is self.tree and device==self.device:return self.plan
|
| 17 |
+
buckets=defaultdict(list);start=tree.prefix_length
|
| 18 |
+
for index,length in enumerate(tree.branch_lengths):
|
| 19 |
+
buckets[((length+31)//32)*32].append((index,start,length));start+=length
|
| 20 |
+
assert start==tree.length
|
| 21 |
+
blocks=[];restore=[None]*len(tree.branch_lengths);offset=0
|
| 22 |
+
for width,branches in sorted(buckets.items()):
|
| 23 |
+
for begin in range(0,len(branches),8):
|
| 24 |
+
block=branches[begin:begin+8];indices=[0]*(8*width);valid=[False]*(8*width)
|
| 25 |
+
for slot,(index,start,length) in enumerate(block):
|
| 26 |
+
a=slot*width;indices[a:a+length]=range(start,start+length);valid[a:a+length]=[True]*length
|
| 27 |
+
restore[index]=torch.arange(length,device=device)+offset+a
|
| 28 |
+
blocks.append((width,torch.tensor(indices,device=device,dtype=torch.long),torch.tensor(valid,device=device,dtype=torch.bool)))
|
| 29 |
+
offset+=8*width
|
| 30 |
+
order=torch.cat(restore) if restore else torch.empty(0,device=device,dtype=torch.long)
|
| 31 |
+
self.tree=tree;self.device=device;self.plan=(blocks,order)
|
| 32 |
+
return self.plan
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def grouped_forward(original,tree_getter,plan_getter=None):
|
| 36 |
+
cache=PackingCache() if plan_getter is None else plan_getter
|
| 37 |
+
def forward(module,value,*args,**kwargs):
|
| 38 |
+
if DEPTH.get():return original(value,*args,**kwargs)
|
| 39 |
+
tree=tree_getter()
|
| 40 |
+
if tree is None:raise ValueError('Grouped token operations require a tree layout')
|
| 41 |
+
if value.ndim not in (2,3) or value.shape[-2]!=tree.length:raise ValueError('Grouped token operations require the packed token sequence')
|
| 42 |
+
if value.ndim==3 and value.shape[0]!=1:raise ValueError('Grouped token operations support one packed request')
|
| 43 |
+
x=value[0] if value.ndim==3 else value
|
| 44 |
+
blocks,order=cache(tree,x.device);token=DEPTH.set(DEPTH.get()+1)
|
| 45 |
+
try:
|
| 46 |
+
prefix=original(x[:tree.prefix_length].contiguous(),*args,**kwargs);outputs=[]
|
| 47 |
+
for width,indices,valid in blocks:
|
| 48 |
+
packed=x.index_select(0,indices).masked_fill(~valid[:,None],0).reshape(8,width,x.shape[-1])
|
| 49 |
+
output=original(packed,*args,**kwargs)
|
| 50 |
+
assert isinstance(output,torch.Tensor) and output.shape[:2]==packed.shape[:2]
|
| 51 |
+
outputs.append(output.reshape(8*width,output.shape[-1]))
|
| 52 |
+
branch=(outputs[0] if len(outputs)==1 else torch.cat(outputs,dim=0)).index_select(0,order) if outputs else prefix[:0]
|
| 53 |
+
output=torch.cat([prefix,branch],dim=0)
|
| 54 |
+
return output[None] if value.ndim==3 else output
|
| 55 |
+
finally:DEPTH.reset(token)
|
| 56 |
+
return forward
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def install(model,*,linears,moe):
|
| 60 |
+
if moe:raise ValueError('MoE dispatch is not a token-local linear operation')
|
| 61 |
+
originals=[]
|
| 62 |
+
def cache(tree, device):
|
| 63 |
+
request=REQUEST.get()
|
| 64 |
+
attention_cache=request[1]
|
| 65 |
+
if not hasattr(attention_cache, "packing"):
|
| 66 |
+
attention_cache.packing=PackingCache()
|
| 67 |
+
return attention_cache.packing(tree,device)
|
| 68 |
+
getter=lambda:REQUEST.get()[0] if REQUEST.get() else None
|
| 69 |
+
if linears:
|
| 70 |
+
for layer in model.backbone.model.layers:
|
| 71 |
+
for module in list(layer.modules()):
|
| 72 |
+
is_lora=all(hasattr(module,name) for name in ['a','b','base','scale'])
|
| 73 |
+
if isinstance(module,torch.nn.Linear) or is_lora:
|
| 74 |
+
original=module.forward;originals.append((module,original))
|
| 75 |
+
module.forward=types.MethodType(grouped_forward(original,getter,cache),module)
|
| 76 |
+
def restore():
|
| 77 |
+
for module,forward in reversed(originals):module.forward=forward
|
| 78 |
+
return restore
|