spragian commited on
Commit
2e0601e
·
verified ·
1 Parent(s): cfcadc3

Release isolated single-pass SPEV decision head

Browse files
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:** a shared scorer over contextual representations of the supplied choices.
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.3234945284875692,
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:e711007c2b401a337442e7b9e1ae6d89c36a2a2dee67cda9a817730b766ca1a1
3
- size 55615444
 
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": 62871149796
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.1.bias": "model-decision.safetensors",
19943
- "head.scorer.1.weight": "model-decision.safetensors",
19944
- "head.scorer.3.bias": "model-decision.safetensors",
19945
- "head.scorer.3.weight": "model-decision.safetensors"
 
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
- "OptionHead",
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 = OptionHead(
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, scoring_positions=None, **kwargs
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
- scoring_positions is None
99
- or scoring_positions.ndim != 2
100
- or scoring_positions.shape[0] != 1
101
- ):
102
- raise ValueError("scoring_positions must have shape (1, candidates)")
103
- if (
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
- hidden = self.backbone(
112
- input_ids=input_ids,
113
- attention_mask=attention_mask,
114
- use_cache=False,
115
- return_dict=True,
116
- **kwargs,
117
- ).logits
118
- positions = scoring_positions.to(hidden.device)
119
- features = hidden.gather(
120
- 1, positions.unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
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=count, temperature=self.config.temperature
 
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
- render,
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, mapping = render({} if state is None else state, questions)
 
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 functools import partial
 
60
 
61
- from transformers.video_utils import load_video
62
-
63
- frames, metadata = load_video(
64
- str(video),
65
- backend="pyav",
66
- sample_indices_fn=partial(
67
- self.video_processor.sample_frames,
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
- positions = locate_markers(ids, self.marker, len(mapping))
99
- inputs["scoring_positions"] = torch.tensor([positions], dtype=torch.long)
 
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