1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
use zeppelin_embed::epoch::EmbeddingTower;
/// Architecture id persisted in `.zem` tower metadata.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u16)]
pub enum Architecture {
/// BERT and XLM-R-compatible post-norm encoder.
Bert = 1,
/// Alibaba GTE `NewModel` encoder with RoPE and gated MLP.
Gte = 2,
/// ArcticQuery: factorized token embeddings, GTE blocks, CLS projection.
ArcticQuery = 3,
}
/// Token pooling policy.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum Pooling {
/// Masked arithmetic mean.
Mean = 1,
/// First token.
Cls = 2,
/// Last non-padding token.
Last = 3,
}
/// Role assigned to one bundle tower.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum TowerRole {
/// Document ingest tower, or the sole symmetric tower.
Document = 0,
/// Query-only tower in an asymmetric pair.
Query = 1,
}
/// Architecture dimensions required to evaluate one tower.
#[derive(Clone, Debug, PartialEq)]
pub struct ArchitectureConfig {
/// Number of transformer blocks.
pub layer_count: u16,
/// Hidden width.
pub hidden: u32,
/// Attention head count.
pub heads: u16,
/// Feed-forward width.
pub intermediate: u32,
/// Token-type vocabulary width.
pub type_vocab: u32,
/// Maximum positional embedding count.
pub max_positions: u32,
/// Layer normalization epsilon.
pub layer_norm_eps: f32,
/// RoPE base for GTE and ArcticQuery.
pub rope_theta: f32,
/// Dense head output width, or zero when absent.
pub dense_out: u32,
/// ArcticQuery token embedding rank; zero for other architectures.
pub embedding_rank: u32,
/// ArcticQuery bias-free output width; zero for other architectures.
pub output_dimension: u32,
}
/// Fully declared model tower loaded from bundle metadata.
#[derive(Clone, Debug, PartialEq)]
pub struct TowerSpec {
/// Slot role.
pub role: TowerRole,
/// Core epoch fields.
pub embedding: EmbeddingTower,
/// Pooling policy.
pub pooling: Pooling,
/// Explicit architecture selector.
pub architecture: Architecture,
/// Architecture dimensions.
pub config: ArchitectureConfig,
}
/// Rectangular padded token batch passed to a runtime.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenBatch {
/// Row-major token ids.
pub token_ids: Vec<i32>,
/// Row-major mask values, one for real tokens and zero for padding.
pub attention_mask: Vec<f32>,
/// Batch rows.
pub rows: usize,
/// Padded tokens per row.
pub tokens_per_row: usize,
}
impl TokenBatch {
/// Re-pads every row to exactly `width` tokens.
///
/// `encode_batch` pads to the longest row it saw, which is what the
/// MLX path wants. CoreML's `embed_batch` pads internally to its fixed
/// sequence length; callers may also request exact padding here. Rows longer
/// than `width` are a caller error, not something to silently trim:
/// truncating a query would change what was asked.
pub fn padded_to(&self, width: usize) -> Result<Self, crate::runtime::RuntimeError> {
if width < self.tokens_per_row {
return Err(crate::runtime::RuntimeError::Shape(format!(
"cannot pad {} tokens per row down to {width}",
self.tokens_per_row
)));
}
let mut token_ids = Vec::with_capacity(self.rows.saturating_mul(width));
let mut attention_mask = Vec::with_capacity(self.rows.saturating_mul(width));
for row in 0..self.rows {
let start = row.saturating_mul(self.tokens_per_row);
let end = start.saturating_add(self.tokens_per_row);
let ids = self.token_ids.get(start..end).ok_or_else(|| {
crate::runtime::RuntimeError::Shape("token id row is out of range".to_owned())
})?;
let mask = self.attention_mask.get(start..end).ok_or_else(|| {
crate::runtime::RuntimeError::Shape("attention row is out of range".to_owned())
})?;
token_ids.extend_from_slice(ids);
attention_mask.extend_from_slice(mask);
token_ids.resize(
token_ids.len().saturating_add(width - self.tokens_per_row),
0,
);
attention_mask.resize(
attention_mask
.len()
.saturating_add(width - self.tokens_per_row),
0.0,
);
}
Self::new(token_ids, attention_mask, self.rows, width)
.map_err(|detail| crate::runtime::RuntimeError::Shape(detail.to_owned()))
}
/// Validates and constructs a rectangular token batch.
pub fn new(
token_ids: Vec<i32>,
attention_mask: Vec<f32>,
rows: usize,
tokens_per_row: usize,
) -> Result<Self, &'static str> {
let expected = rows
.checked_mul(tokens_per_row)
.ok_or("token batch shape overflow")?;
if rows == 0 || tokens_per_row == 0 || token_ids.len() != expected {
return Err("token ids do not match the declared batch shape");
}
if attention_mask.len() != expected {
return Err("attention mask does not match the declared batch shape");
}
Ok(Self {
token_ids,
attention_mask,
rows,
tokens_per_row,
})
}
}