diff --git a/.gitignore b/.gitignore index 58bdb0970..6e456c177 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ tokenizer.json out.txt todo.md rust-analyzer +*.gguf /.claude **/.claude/ profiles diff --git a/Cargo.lock b/Cargo.lock index 79ac05f6b..49ac032d3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -217,7 +217,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -228,7 +228,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2032,7 +2032,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.2", ] [[package]] @@ -2053,7 +2053,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -4048,7 +4048,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4404,7 +4404,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6203,7 +6203,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi 0.5.2", "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6374,6 +6374,14 @@ dependencies = [ "tracing", ] +[[package]] +name = "kalosm-chunking" +version = "0.4.0" +dependencies = [ + "srx", + "whatlang", +] + [[package]] name = "kalosm-common" version = "0.4.0" @@ -6430,6 +6438,7 @@ dependencies = [ "heed", "image 0.24.9", "kalosm", + "kalosm-chunking", "kalosm-language-model", "kalosm-llama", "kalosm-sample", @@ -6445,7 +6454,6 @@ dependencies = [ "scraper", "serde", "serde_json", - "srx", "surrealdb", "tempfile", "texting_robots", @@ -7611,7 +7619,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9362,6 +9370,24 @@ version = "0.8.53" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +[[package]] +name = "rgliner" +version = "0.4.0" +dependencies = [ + "anyhow", + "clap 4.6.1", + "fusor", + "kalosm-common", + "kalosm-language-model", + "kalosm-model-types", + "rbert", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokenizers 0.21.4", + "tokio", +] + [[package]] name = "ring" version = "0.17.14" @@ -9622,7 +9648,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -10302,7 +10328,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -10403,7 +10429,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -10932,7 +10958,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -12495,7 +12521,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 3c868ee93..73fd96894 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ bench = false [workspace] members = [ "models/rbert", + "models/rgliner", "models/kalosm-llama", "models/kalosm-tokenizer", "models/rwhisper", @@ -29,6 +30,7 @@ members = [ "interfaces/kalosm-streams", "interfaces/kalosm-parse-macro", "interfaces/kalosm-common", + "interfaces/kalosm-chunking", "interfaces/kalosm-model-types", "fusor-ml/core", "fusor-ml/gguf", @@ -51,6 +53,7 @@ kalosm = { path = "./interfaces/kalosm", version = "0.4.0", default-features = f kalosm-sample = { path = "./interfaces/kalosm-sample", version = "0.4.0" } kalosm-parse-macro = { path = "./interfaces/kalosm-parse-macro", version = "0.4.0" } kalosm-common = { path = "./interfaces/kalosm-common", version = "0.4.0" } +kalosm-chunking = { path = "./interfaces/kalosm-chunking", version = "0.4.0" } kalosm-model-types = { path = "./interfaces/kalosm-model-types", version = "0.4.0" } kalosm-language-model = { path = "./interfaces/language-model", version = "0.4.0" } kalosm-streams = { path = "./interfaces/kalosm-streams", version = "0.4.0" } @@ -58,6 +61,7 @@ kalosm-language = { path = "./interfaces/kalosm-language", version = "0.4.0", de kalosm-sound = { path = "./interfaces/kalosm-sound", version = "0.4.0", default-features = false } kalosm-vision = { path = "./interfaces/kalosm-vision", version = "0.4.0", default-features = false } rbert = { path = "./models/rbert", version = "0.4.0", default-features = false } +rgliner = { path = "./models/rgliner", version = "0.4.0" } kalosm-llama = { path = "./models/kalosm-llama", version = "0.4.0", default-features = false } kalosm-tokenizer = { path = "./models/kalosm-tokenizer", version = "0.4.0" } rwhisper = { path = "./models/rwhisper", version = "0.4.0", default-features = false } diff --git a/fusor-ml/core/src/compute_graph/resolve/run.rs b/fusor-ml/core/src/compute_graph/resolve/run.rs index 12cc100c3..1bd22c7b3 100644 --- a/fusor-ml/core/src/compute_graph/resolve/run.rs +++ b/fusor-ml/core/src/compute_graph/resolve/run.rs @@ -533,7 +533,18 @@ impl Resolver { let encode_start = host_trace.then(Instant::now); let mut dispatch_index = 0usize; let mut command_index = 0usize; - let dispatches_per_pass = dispatches_per_pass(total_kernels); + // Bound the cumulative workgroups recorded into a single compute pass by + // the device's real per-dimension workgroup limit. Each + // `begin_compute_pass` is a GPU command-encoder; recording too many + // passes into one command buffer loses the device on Metal (a + // long-sequence encoder + BiLSTM resolves to ~2200 dispatches, and one + // pass per dispatch — the old policy for large graphs — meant ~2200 + // passes, which faults; the same dispatches grouped into a handful of + // passes run fine). Grouping by the device-reported workgroup limit + // keeps the pass count at roughly `total_workgroups / limit`: a few + // dozen for the largest graphs, one for decode/prefill — no threshold. + let pass_workgroup_budget = + u64::from(device.limits().max_compute_workgroups_per_dimension); while command_index < commands.len() { match &commands[command_index] { CommandRecord::CopyBuffer(copy) => { @@ -577,14 +588,20 @@ impl Resolver { label: Some("Resolver Direct Kernels"), timestamp_writes: None, }); - let mut pass_dispatches = 0usize; + let mut pass_workgroups = 0u64; while command_index < commands.len() { - if pass_dispatches >= dispatches_per_pass { - break; - } let CommandRecord::Dispatch(record) = &commands[command_index] else { break; }; + let dispatch_workgroups = record.dispatch.workgroup_count(); + // Always record at least one dispatch per pass; otherwise + // close the pass before its cumulative workgroups would + // exceed the device budget. + if pass_workgroups > 0 + && pass_workgroups + dispatch_workgroups > pass_workgroup_budget + { + break; + } if let Some((query_set, _, _, _)) = &query_resources { pass.write_timestamp(query_set, (dispatch_index * 2) as u32); } @@ -596,7 +613,7 @@ impl Resolver { } dispatch_index += 1; command_index += 1; - pass_dispatches += 1; + pass_workgroups += dispatch_workgroups; } } } @@ -715,14 +732,3 @@ fn direct_plan_binding_buffers(inputs: &[MirValue]) -> Vec usize { - if let Ok(value) = std::env::var("FUSOR_RESOLVE_DISPATCHES_PER_PASS") - && let Ok(parsed) = value.parse::() - && parsed > 0 - { - return parsed; - } - - if total_kernels >= 1024 { 1 } else { usize::MAX } -} diff --git a/fusor-ml/cpu/benches/comparison.rs b/fusor-ml/cpu/benches/comparison.rs index 299b25e31..977bf433f 100644 --- a/fusor-ml/cpu/benches/comparison.rs +++ b/fusor-ml/cpu/benches/comparison.rs @@ -1,6 +1,7 @@ use candle_core::{Device, Tensor as CandleTensor}; -use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use fusor_cpu::__private::TypedTensor; +use std::hint::black_box; const SIZES: &[usize] = &[64, 256]; diff --git a/fusor-ml/cpu/benches/fusion.rs b/fusor-ml/cpu/benches/fusion.rs index 73c4d7776..016c4599b 100644 --- a/fusor-ml/cpu/benches/fusion.rs +++ b/fusor-ml/cpu/benches/fusion.rs @@ -5,8 +5,9 @@ //! memory multiple times (once per operation). use candle_core::{Device, Tensor as CandleTensor}; -use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use fusor_cpu::__private::{ConcreteTensor, TypedTensor}; +use std::hint::black_box; /// Benchmark fused operations (single memory pass) vs separate operations (multiple passes) fn bench_fusion(c: &mut Criterion) { diff --git a/fusor-ml/cpu/benches/qmatmul.rs b/fusor-ml/cpu/benches/qmatmul.rs index a7b3980e8..d20d9aee6 100644 --- a/fusor-ml/cpu/benches/qmatmul.rs +++ b/fusor-ml/cpu/benches/qmatmul.rs @@ -1,8 +1,9 @@ use aligned_vec::AVec; -use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use fusor_cpu::__private::{ConcreteTensor, QuantizedTensor, TypedTensor}; use fusor_cpu::BlockQ8_0; use half::f16; +use std::hint::black_box; /// Helper to create a Q8_0 block from scale and data fn make_q8_0_block(scale: f32, data: [i8; 32]) -> BlockQ8_0 { diff --git a/fusor-ml/cpu/benches/tensor_ops.rs b/fusor-ml/cpu/benches/tensor_ops.rs index 91890c8e9..4f6782483 100644 --- a/fusor-ml/cpu/benches/tensor_ops.rs +++ b/fusor-ml/cpu/benches/tensor_ops.rs @@ -1,5 +1,6 @@ -use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use fusor_cpu::__private::{Abs, Add, ConcreteTensor, Div, Mul, Neg, Sqrt, Sub, TensorBacking}; +use std::hint::black_box; const SIZES: &[usize] = &[64, 256]; diff --git a/fusor-ml/fusor/src/composite/mod.rs b/fusor-ml/fusor/src/composite/mod.rs index 36cb1dcba..b529f8d1f 100644 --- a/fusor-ml/fusor/src/composite/mod.rs +++ b/fusor-ml/fusor/src/composite/mod.rs @@ -17,6 +17,7 @@ mod reductions; mod rope; mod shape; mod to_vec; +mod transformer; mod upsample; mod where_cond; @@ -25,3 +26,8 @@ pub use rope::{RopeCache, base_inverse_frequency}; pub(crate) use shape::broadcast_shapes; pub use shape::{arange, arange_step, cat, stack}; pub use to_vec::{ToVec, ToVec1, ToVec2, ToVec3}; +pub use transformer::{ + AttentionBias, AttentionVariant, FeedForwardVariant, GatedActivation, GroupedAttention, + LlamaFeedForward, Norm, PhiFeedForward, RopeLike, SeparateAttention, TransformerBlock, + forward_attention_qkv_f32, +}; diff --git a/fusor-ml/fusor/src/composite/transformer.rs b/fusor-ml/fusor/src/composite/transformer.rs new file mode 100644 index 000000000..83a171bd3 --- /dev/null +++ b/fusor-ml/fusor/src/composite/transformer.rs @@ -0,0 +1,921 @@ +//! Shared pre-norm transformer block. +//! +//! A single parameterized attention + feed-forward block used by every +//! transformer model in the workspace (the Llama/Qwen decoders, the Qwen +//! vision tower, and the ModernBERT / Qwen encoders). The block composes the +//! low-level fusor primitives (`flash_attention`, RoPE, `RmsNorm`/`LayerNorm`, +//! `KvCache`, quantized `Linear`) behind a small set of enums so the same code +//! serves causal decoders and bidirectional encoders alike: +//! +//! * [`AttentionVariant`] — fused (`Grouped`) vs separate (`Separate`) Q/K/V +//! projection, with optional grouped-query attention and optional q/k norm. +//! * [`FeedForwardVariant`] — gated SwiGLU/GeGLU ([`LlamaFeedForward`]) vs the +//! plain split-gate Phi MLP ([`PhiFeedForward`]); the gated activation is +//! selectable via [`GatedActivation`]. +//! * [`Norm`] — `RmsNorm` or `LayerNorm`, so decoders (RMS) and BERT-style +//! encoders (Layer) share the same block. +//! * [`RopeLike`] — abstracts the rotary cache so the block is agnostic to the +//! concrete RoPE implementation (plain [`RopeCache`] or a model-specific one). +//! +//! The decoder hot path uses [`TransformerBlock::forward`] (attention sublayer +//! only; the model loop orchestrates norms + residuals so it can fuse them). +//! Encoders use [`TransformerBlock::forward_block`], which runs the full +//! pre-norm block in one call. + +use crate::D; +use crate::MaskKind; +use crate::QMatrix; +use crate::RopeCache; +use crate::Tensor; +use crate::cache::AttentionMask; +use crate::cache::KvCache; +use crate::layers::{LayerNorm, Linear, RmsNorm}; +use crate::{CastTensor, CastTo, FloatDataType, Fusion, SimdElement}; + +/// Abstracts a rotary-position-embedding cache so [`TransformerBlock`] does not +/// depend on any one model's RoPE implementation. Implemented for the plain +/// [`RopeCache`]; model-specific caches (e.g. multi-axis vision RoPE) provide +/// their own impl. +pub trait RopeLike { + /// Apply RoPE to the query and key tensors, returning the rotated pair. + fn apply( + &self, + query: &Tensor<4, F>, + key: &Tensor<4, F>, + start_pos: usize, + position_ids: Option<&Tensor<2, F>>, + interleaved: bool, + ) -> (Tensor<4, F>, Tensor<4, F>); +} + +impl RopeLike for RopeCache +where + F: FloatDataType + SimdElement + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn apply( + &self, + query: &Tensor<4, F>, + key: &Tensor<4, F>, + start_pos: usize, + _position_ids: Option<&Tensor<2, F>>, + interleaved: bool, + ) -> (Tensor<4, F>, Tensor<4, F>) { + let q_f32: Tensor<4, f32> = query.cast(); + let k_f32: Tensor<4, f32> = key.cast(); + let (q_out, k_out) = if interleaved { + self.forward_interleaved(&q_f32, &k_f32, start_pos) + } else { + self.forward(&q_f32, &k_f32, start_pos) + }; + (q_out.cast(), k_out.cast()) + } +} + +/// Activation applied to the gate branch of a gated feed-forward network. +#[derive(Clone, Copy, Debug)] +pub enum GatedActivation { + /// SiLU / Swish — used by Llama, Qwen, Gemma (SwiGLU). + SiLU, + /// GELU — used by ModernBERT (GeGLU). + GeLU, +} + +impl GatedActivation { + #[inline] + fn apply(self, x: &Tensor<3, f32>) -> Tensor<3, f32> { + match self { + GatedActivation::SiLU => x.silu(), + GatedActivation::GeLU => x.gelu().to_concrete(), + } + } +} + +/// Normalization layer used inside a [`TransformerBlock`]. `RmsNorm` for +/// decoders / Qwen encoders, `LayerNorm` (f32) for BERT-style encoders. +pub enum Norm { + /// Root-mean-square normalization (no centering). + Rms(RmsNorm<1, F>), + /// Standard layer normalization (computed in f32). + Layer(LayerNorm<1, f32>), +} + +impl Norm +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + /// Pre-normalization of a hidden state. + pub fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + match self { + Norm::Rms(n) => n.forward_generic(x), + Norm::Layer(n) => { + let x_f32 = x.cast::(); + let out: Tensor<3, f32> = n.forward(&x_f32).to_concrete(); + out.cast() + } + } + } + + /// Fused `(input + residual)` followed by normalization. The `Rms` arm uses + /// the fused residual kernel (the decode hot path depends on this); the + /// `Layer` arm falls back to an explicit add (no `LayerNorm` fused-residual + /// kernel exists, and encoders never relied on one). + pub fn forward_residual_f32( + &self, + input: &Tensor<3, f32, B1>, + residual: &Tensor<3, f32, B2>, + ) -> Tensor<3, F> + where + B1: Fusion<3, f32>, + B2: Fusion<3, f32>, + { + match self { + Norm::Rms(n) => n.forward_residual_f32(input, residual), + Norm::Layer(n) => { + let sum = input.add_(residual); + let out: Tensor<3, f32> = n.forward(&sum).to_concrete(); + out.cast() + } + } + } +} + +/// Gated vs plain feed-forward selection. +pub enum FeedForwardVariant { + /// Gated FFN (SwiGLU / GeGLU) — Llama, Qwen, Gemma, ModernBERT. + Llama(Box>), + /// Plain split-gate FFN — Phi models. + Phi(PhiFeedForward), +} + +impl FeedForwardVariant +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + /// Feed-forward forward pass. + pub fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + match self { + FeedForwardVariant::Llama(ffn) => ffn.forward(x), + FeedForwardVariant::Phi(ffn) => ffn.forward(x), + } + } + + /// Feed-forward pass that folds two residual adds into the down-projection + /// epilogue (decode fast path). Returns `None` when the fold does not apply. + pub fn forward_add_residuals( + &self, + x: &Tensor<3, F, B>, + first: &Tensor<3, f32, B1>, + second: &Tensor<3, f32, B2>, + ) -> Option> + where + B: Fusion<3, F>, + B1: Fusion<3, f32>, + B2: Fusion<3, f32>, + { + match self { + FeedForwardVariant::Llama(ffn) => ffn.forward_add_residuals(x, first, second), + FeedForwardVariant::Phi(_) => None, + } + } +} + +/// Plain split-gate feed-forward used by Phi models. +pub struct PhiFeedForward { + /// Up projection (its output is split into gate + up halves). + pub up: QMatrix, + /// Down projection. + pub down: QMatrix, + /// Intermediate width (half of the up projection output). + pub feed_forward_length: usize, +} + +impl PhiFeedForward { + fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + { + // All computation happens in f32 for compatibility with SIMD ops + let x_f32 = x.cast::(); + let up_states = x_f32.q_mat_mul(&self.up); + let gate = up_states + .narrow(D::Minus1, 0, self.feed_forward_length) + .to_concrete(); + let up_states = up_states + .narrow( + D::Minus1, + self.feed_forward_length, + self.feed_forward_length, + ) + .to_concrete(); + let gate = gate.silu(); + let up_states = up_states * gate; + let result = up_states.q_mat_mul(&self.down); + result.cast() + } +} + +/// Gated feed-forward (SwiGLU / GeGLU). Supports a pre-fused `gate_up` +/// projection (single matmul) or separate gate/up projections, with optional +/// per-projection biases, and a selectable [`GatedActivation`]. +pub struct LlamaFeedForward { + gate: Option, + gate_up: Option, + gate_bias: Option>, + down: QMatrix, + down_bias: Option>, + up: Option, + up_bias: Option>, + activation: GatedActivation, +} + +impl LlamaFeedForward { + /// Gated FFN from separate gate/up/down projections (SiLU activation). + pub fn new(gate: QMatrix, down: QMatrix, up: QMatrix) -> Self { + let gate_up = QMatrix::concat_rows(&[&gate, &up]); + Self { + gate: Some(gate), + gate_up, + down, + up: Some(up), + gate_bias: None, + down_bias: None, + up_bias: None, + activation: GatedActivation::SiLU, + } + } + + /// Gated FFN from a pre-fused `[2 * intermediate, hidden]` gate+up + /// projection and a down projection, with a selectable activation. Used by + /// ModernBERT (GeGLU) where the gate and up weights are stored fused. + pub fn from_fused_gated(gate_up: QMatrix, down: QMatrix, activation: GatedActivation) -> Self { + Self { + gate: None, + gate_up: Some(gate_up), + down, + up: None, + gate_bias: None, + down_bias: None, + up_bias: None, + activation, + } + } + + /// Gated FFN with optional per-projection biases (SiLU activation). + pub fn new_with_bias( + gate: QMatrix, + gate_bias: Option>, + down: QMatrix, + down_bias: Option>, + up: QMatrix, + up_bias: Option>, + ) -> Self { + let gate_up = QMatrix::concat_rows(&[&gate, &up]); + Self { + gate: Some(gate), + gate_up, + gate_bias, + down, + down_bias, + up: Some(up), + up_bias, + activation: GatedActivation::SiLU, + } + } + + /// Gated feed-forward pass: `down(act(gate(x)) * up(x))`. + pub fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> + where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + { + let up_result = self.activation(x); + let mut up = up_result.q_mat_mul(&self.down); + if let Some(ref bias) = self.down_bias { + let bias_f32: Tensor<1, f32> = bias.cast(); + up = up.add_(&bias_f32); + } + + // Cast back to F + up.cast() + } + + fn forward_add_residuals( + &self, + x: &Tensor<3, F, B>, + first: &Tensor<3, f32, B1>, + second: &Tensor<3, f32, B2>, + ) -> Option> + where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + B1: Fusion<3, f32>, + B2: Fusion<3, f32>, + { + if self.down_bias.is_some() { + return None; + } + if x.shape()[1] > 1 { + return None; + } + + let up_result = self.activation(x); + // Residual adds authored in natural graph form: the resolver folds both + // `add`s into the qmatmul post epilogue (one dispatch on decode). + let projected = up_result.q_mat_mul(&self.down); + let with_first = (&projected + first).to_concrete(); + let up = (&with_first + second).to_concrete(); + Some(up.cast()) + } + + fn activation(&self, x: &Tensor<3, F, B>) -> Tensor<3, f32> + where + F: CastTo + CastTensor, + B: Fusion<3, F>, + { + // All computation happens in f32 for compatibility with SIMD ops + let x_f32 = x.cast::(); + + match &self.gate_up { + Some(gate_up) if self.gate_bias.is_none() && self.up_bias.is_none() => { + // SwiGLU/GeGLU split/gate authored in natural graph form: the + // resolver folds `act(gate) * up` over the two narrow halves into + // the qmatmul accumulator-offset epilogue (one dispatch on decode). + let pair_len = gate_up.shape()[0] / 2; + let projected = x_f32.q_mat_mul(gate_up); + let gate = projected.narrow(D::Minus1, 0, pair_len).to_concrete(); + let up = projected + .narrow(D::Minus1, pair_len, pair_len) + .to_concrete(); + (self.activation.apply(&gate) * up).to_concrete() + } + Some(gate_up) => { + let gate_width = self.gate.as_ref().expect("gated ffn gate").shape()[0]; + let up_width = self.up.as_ref().expect("gated ffn up").shape()[0]; + let gate_up_states = x_f32.q_mat_mul(gate_up); + + let mut gate_states = gate_up_states + .narrow(D::Minus1, 0, gate_width) + .to_concrete(); + if let Some(ref bias) = self.gate_bias { + let bias_f32: Tensor<1, f32> = bias.cast(); + gate_states = gate_states.add_(&bias_f32); + } + + let mut up_states = gate_up_states + .narrow(D::Minus1, gate_width, up_width) + .to_concrete(); + if let Some(ref bias) = self.up_bias { + let bias_f32: Tensor<1, f32> = bias.cast(); + up_states = up_states.add_(&bias_f32); + } + + (self.activation.apply(&gate_states) * up_states).to_concrete() + } + None => { + let gate = self.gate.as_ref().expect("gated ffn gate"); + let up = self.up.as_ref().expect("gated ffn up"); + let mut w1 = x_f32.q_mat_mul(gate); + if let Some(ref bias) = self.gate_bias { + let bias_f32: Tensor<1, f32> = bias.cast(); + w1 = w1.add_(&bias_f32); + } + let w1 = self.activation.apply(&w1); + + let mut w3 = x_f32.q_mat_mul(up); + if let Some(ref bias) = self.up_bias { + let bias_f32: Tensor<1, f32> = bias.cast(); + w3 = w3.add_(&bias_f32); + } + + (w1 * w3).to_concrete() + } + } + } +} + +/// Fused vs separate Q/K/V projection. +pub enum AttentionVariant { + /// Separate Q/K/V projections (optionally a fused weight + optional q/k norm). + Separate(Box>), + /// Single fused Q/K/V projection. + Grouped(GroupedAttention), +} + +impl AttentionVariant +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + /// Project + RoPE the hidden states into `(query, key, value)` head tensors. + #[allow(clippy::too_many_arguments)] + pub fn forward( + &self, + num_heads: usize, + head_dim: usize, + num_key_value_heads: usize, + hidden_states: &Tensor<3, F, B>, + rope: &R, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + ) -> (Tensor<4, F>, Tensor<4, F>, Tensor<4, F>) + where + B: Fusion<3, F>, + R: RopeLike, + { + match self { + AttentionVariant::Separate(attention) => attention.forward( + num_heads, + head_dim, + num_key_value_heads, + hidden_states, + rope, + start_pos, + pos_ids, + ), + AttentionVariant::Grouped(attention) => attention.forward( + num_heads, + head_dim, + num_key_value_heads, + hidden_states, + rope, + start_pos, + pos_ids, + ), + } + } +} + +/// Optional additive biases for the Q/K/V projections. +pub struct AttentionBias { + bias_q: Tensor<1, F>, + bias_k: Tensor<1, F>, + bias_v: Tensor<1, F>, + bias_qkv: Tensor<1, F>, +} + +impl AttentionBias { + /// Build an attention bias from separate q/k/v biases (also concatenated + /// for the fused-QKV path). + pub fn new(q: Tensor<1, F>, k: Tensor<1, F>, v: Tensor<1, F>) -> Self { + let bias_qkv = crate::cat([q.clone(), k.clone(), v.clone()], 0).to_concrete(); + Self { + bias_q: q, + bias_k: k, + bias_v: v, + bias_qkv, + } + } +} + +/// Separate Q/K/V projection (optionally backed by a fused weight, optionally +/// with per-head q/k normalization, optionally with biases). +pub struct SeparateAttention { + /// Query projection. + pub attention_wq: QMatrix, + /// Optional fused Q/K/V projection (used when present). + pub attention_qkv: Option, + /// Optional per-head query normalization. + pub attention_q_norm: Option>, + /// Key projection. + pub attention_wk: QMatrix, + /// Optional per-head key normalization. + pub attention_k_norm: Option>, + /// Value projection. + pub attention_wv: QMatrix, + /// Optional projection biases. + pub bias: Option>, + /// Whether RoPE pairs adjacent elements (interleaved) or halves. + pub interleaved_rope: bool, +} + +impl SeparateAttention +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + num_heads: usize, + head_dim: usize, + num_key_value_heads: usize, + hidden_states: &Tensor<3, F, B>, + rope: &R, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + ) -> (Tensor<4, F>, Tensor<4, F>, Tensor<4, F>) + where + B: Fusion<3, F>, + R: RopeLike, + { + let [b_sz, seq_len, _] = hidden_states.shape(); + + // Compute in f32 for SIMD ops compatibility + let hidden_f32 = hidden_states.cast::(); + + if let Some(attention_qkv) = &self.attention_qkv { + let query_width = num_heads * head_dim; + let key_width = num_key_value_heads * head_dim; + let value_width = key_width; + let mut qkv = hidden_f32.q_mat_mul(attention_qkv); + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_qkv.cast(); + qkv = qkv.add_(&bias_f32); + } + + let query_states: Tensor<4, F> = { + let query_states = qkv.narrow(D::Minus1, 0, query_width).to_concrete(); + + let query = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + + let query: Tensor<4, F> = query.cast(); + if let Some(norm) = &self.attention_q_norm { + norm.forward_generic_4d(&query) + } else { + query + } + }; + + let key_states: Tensor<4, F> = { + let key_states = qkv.narrow(D::Minus1, query_width, key_width).to_concrete(); + + let key = key_states + .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + + let key: Tensor<4, F> = key.cast(); + if let Some(norm) = &self.attention_k_norm { + norm.forward_generic_4d(&key) + } else { + key + } + }; + + let value_states: Tensor<4, F> = { + let value_states = qkv + .narrow(D::Minus1, query_width + key_width, value_width) + .to_concrete(); + + value_states + .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) + .transpose(1, 2) + .to_concrete() + .cast() + }; + + let (query_states, key_states) = rope.apply( + &query_states, + &key_states, + start_pos, + pos_ids, + self.interleaved_rope, + ); + return (query_states, key_states, value_states); + } + + let query_states: Tensor<4, F> = { + let mut query_states = hidden_f32.q_mat_mul(&self.attention_wq); + + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_q.cast(); + query_states = query_states.add_(&bias_f32); + } + + let query = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + + let query: Tensor<4, F> = query.cast(); + if let Some(norm) = &self.attention_q_norm { + norm.forward_generic_4d(&query) + } else { + query + } + }; + let key_states: Tensor<4, F> = { + let mut key_states = hidden_f32.q_mat_mul(&self.attention_wk); + + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_k.cast(); + key_states = key_states.add_(&bias_f32); + } + + let key = key_states + .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + + let key: Tensor<4, F> = key.cast(); + if let Some(norm) = &self.attention_k_norm { + norm.forward_generic_4d(&key) + } else { + key + } + }; + let value_states: Tensor<4, F> = { + let mut value_states = hidden_f32.q_mat_mul(&self.attention_wv); + + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_v.cast(); + value_states = value_states.add_(&bias_f32); + } + + value_states + .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) + .transpose(1, 2) + .to_concrete() + .cast() + }; + + let (query_states, key_states) = rope.apply( + &query_states, + &key_states, + start_pos, + pos_ids, + self.interleaved_rope, + ); + (query_states, key_states, value_states) + } +} + +/// Single fused Q/K/V projection (grouped-query friendly). +pub struct GroupedAttention { + /// Fused `[(q + k + v) heads * head_dim, hidden]` projection. + pub attention_qkv: QMatrix, + /// Whether RoPE pairs adjacent elements (interleaved) or halves. + pub interleaved_rope: bool, +} + +impl GroupedAttention { + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + num_heads: usize, + head_dim: usize, + num_key_value_heads: usize, + x: &Tensor<3, F, B>, + rope: &R, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + ) -> (Tensor<4, F>, Tensor<4, F>, Tensor<4, F>) + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + R: RopeLike, + { + let [b_sz, seq_len, _] = x.shape(); + // Compute in f32 for SIMD ops compatibility + let x_f32 = x.cast::(); + let qkv = x_f32.q_mat_mul(&self.attention_qkv); + + let query_pos = num_heads * head_dim; + let query_states = qkv.narrow(D::Minus1, 0, query_pos); + let key_states = qkv.narrow(D::Minus1, query_pos, num_key_value_heads * head_dim); + let value_states = qkv.narrow( + D::Minus1, + query_pos + num_key_value_heads * head_dim, + num_key_value_heads * head_dim, + ); + + let query_states: Tensor<4, F> = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete() + .cast(); + let key_states: Tensor<4, F> = key_states + .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) + .transpose(1, 2) + .to_concrete() + .cast(); + let value_states: Tensor<4, F> = value_states + .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) + .transpose(1, 2) + .to_concrete() + .cast(); + + let (query_states, key_states) = rope.apply( + &query_states, + &key_states, + start_pos, + pos_ids, + self.interleaved_rope, + ); + + (query_states, key_states, value_states) + } +} + +/// A complete pre-norm transformer block: attention sublayer + gated/plain FFN, +/// each with its own normalization and (optional) post-normalization. +/// +/// Generic over the float type `F` (decoders may run f16; encoders are f32) and +/// the rotary cache `R` (so a vision tower can supply multi-axis RoPE). +pub struct TransformerBlock { + /// Q/K/V projection variant. + pub attention_variant: AttentionVariant, + /// Output projection. + pub attention_wo: Linear, + /// Pre-attention normalization (`None` only for blocks whose input is + /// pre-normalized upstream, e.g. ModernBERT layer 0). + pub attention_norm: Option>, + /// Optional post-attention normalization. + pub post_attention_norm: Option>, + /// Feed-forward variant. + pub feed_forward_variant: FeedForwardVariant, + /// Pre-FFN normalization. + pub ffn_norm: Norm, + /// Optional post-FFN normalization. + pub post_ffn_norm: Option>, + /// Number of query heads. + pub n_head: usize, + /// Number of key/value heads (== `n_head` for MHA, fewer for GQA). + pub n_kv_head: usize, + /// Per-head dimension. + pub head_dim: usize, + /// Model hidden size. + pub hidden_size: usize, + /// Rotary cache. + pub rope_cache: R, + /// Sliding-window size for local-attention decoder layers. + pub sliding_window_size: Option, +} + +impl TransformerBlock +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, + R: RopeLike, +{ + /// Attention sublayer only (Q/K/V projection + RoPE + KV cache + flash + /// attention + output projection). The decoder model loop applies the + /// surrounding norms and residuals so it can fuse them; the hidden state + /// passed in is already pre-normalized. + pub fn forward( + &self, + hidden_states: &Tensor<3, F, B>, + attention_mask: Option<&AttentionMask>, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + cache: Option<&mut KvCache>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let [b_sz, q_len, _] = hidden_states.shape(); + let hidden_size = self.hidden_size; + + let (query_states, key_states, value_states) = self.attention_variant.forward( + self.n_head, + self.head_dim, + self.n_kv_head, + hidden_states, + &self.rope_cache, + start_pos, + pos_ids, + ); + + // Convert to f32 for cache operations (cache uses f32 for SIMD compatibility) + let query_f32: Tensor<4, f32> = query_states.cast(); + let key_f32: Tensor<4, f32> = key_states.cast(); + let value_f32: Tensor<4, f32> = value_states.cast(); + + let (key_f32, value_f32) = match cache { + None => (key_f32, value_f32), + Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), + }; + + let mask = attention_mask.map(causal_mask_tuple); + + forward_attention_qkv_f32( + &query_f32, + &key_f32, + &value_f32, + &self.attention_wo, + mask, + self.head_dim, + b_sz, + q_len, + hidden_size, + ) + } + + /// Full pre-norm block in a single call (used by bidirectional encoders): + /// `norm → attention → residual → norm → FFN → residual`, with a + /// `[batch, key]` padding bias applied as a [`MaskKind::BatchKeyMask`]. No + /// KV cache; positions start at 0. + pub fn forward_block( + &self, + hidden_states: &Tensor<3, F>, + mask_bias: Option<&Tensor<2, f32>>, + ) -> Tensor<3, F> + where + crate::AddOp: crate::SimdBinaryOp, + { + let attn = self.attention_sublayer(hidden_states, mask_bias); + let hidden = hidden_states.add_(&attn); + + let ffn_input = self.ffn_norm.forward(&hidden); + let ffn_output = self.feed_forward_variant.forward(&ffn_input); + hidden.add_(&ffn_output) + } + + /// Pre-norm attention sublayer (`norm → attention → output projection`), + /// without the residual add. Exposed so encoders with bespoke masking + /// (e.g. ModernBERT's sliding window) can reuse the shared projection + + /// RoPE while supplying their own attention computation around it. + pub fn attention_sublayer( + &self, + hidden_states: &Tensor<3, F>, + mask_bias: Option<&Tensor<2, f32>>, + ) -> Tensor<3, F> { + let [b_sz, seq_len, _] = hidden_states.shape(); + let normed = match &self.attention_norm { + Some(n) => n.forward(hidden_states), + None => hidden_states.clone(), + }; + let (query_states, key_states, value_states) = self.attention_variant.forward( + self.n_head, + self.head_dim, + self.n_kv_head, + &normed, + &self.rope_cache, + 0, + None, + ); + let query_f32: Tensor<4, f32> = query_states.cast(); + let key_f32: Tensor<4, f32> = key_states.cast(); + let value_f32: Tensor<4, f32> = value_states.cast(); + let mask = mask_bias.map(|m| (m, MaskKind::BatchKeyMask)); + forward_attention_qkv_f32( + &query_f32, + &key_f32, + &value_f32, + &self.attention_wo, + mask, + self.head_dim, + b_sz, + seq_len, + self.hidden_size, + ) + } +} + +/// Map a high-level [`AttentionMask`] to the `(mask, kind)` tuple consumed by +/// `flash_attention`: strictly-causal masks use the GPU-optimized causal kernel, +/// others fall back to an explicit Q×K additive mask. +fn causal_mask_tuple(m: &AttentionMask) -> (&Tensor<2, f32>, MaskKind) { + let kind = if m.is_strict_causal() { + MaskKind::Causal + } else { + MaskKind::QKMask + }; + (m.mask(), kind) +} + +/// Flash attention over `[batch, heads, seq, head_dim]` q/k/v followed by the +/// output projection. Computation is in f32 (SIMD compatibility); the result is +/// cast back to `F`. +#[allow(clippy::too_many_arguments)] +pub fn forward_attention_qkv_f32( + query_states: &Tensor<4, f32>, + key_states: &Tensor<4, f32>, + value_states: &Tensor<4, f32>, + attention_wo: &Linear, + mask: Option<(&Tensor<2, f32>, MaskKind)>, + head_dim: usize, + b_sz: usize, + q_len: usize, + hidden_size: usize, +) -> Tensor<3, F> +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + let scale = 1. / (head_dim as f64).sqrt(); + let attn_output = query_states.flash_attention(key_states, value_states, scale as f32, mask); + + let attn_output = attn_output.transpose(1, 2); + + let attn_output = attn_output.reshape([b_sz, q_len, hidden_size]); + + attention_wo.forward_generic(&attn_output.cast()) +} diff --git a/fusor-ml/fusor/src/lib.rs b/fusor-ml/fusor/src/lib.rs index 1d3360641..21b35f1d1 100644 --- a/fusor-ml/fusor/src/lib.rs +++ b/fusor-ml/fusor/src/lib.rs @@ -29,8 +29,10 @@ pub use quantized::{CpuF32Tensor, QMatrix}; use std::ops::{Deref, Range}; pub use composite::{ - MaskKind, RopeCache, ToVec, ToVec1, ToVec2, ToVec3, arange, arange_step, - base_inverse_frequency, cat, stack, + AttentionBias, AttentionVariant, FeedForwardVariant, GatedActivation, GroupedAttention, + LlamaFeedForward, MaskKind, Norm, PhiFeedForward, RopeCache, RopeLike, SeparateAttention, + ToVec, ToVec1, ToVec2, ToVec3, TransformerBlock, arange, arange_step, base_inverse_frequency, + cat, forward_attention_qkv_f32, stack, }; pub use device::Device; pub use error::Error; diff --git a/fusor-ml/tile-ir-runtime/src/direct_kernel.rs b/fusor-ml/tile-ir-runtime/src/direct_kernel.rs index fa6349ade..b05959e54 100644 --- a/fusor-ml/tile-ir-runtime/src/direct_kernel.rs +++ b/fusor-ml/tile-ir-runtime/src/direct_kernel.rs @@ -538,6 +538,19 @@ impl PreparedDirectDispatch { self.steps.len() } + /// Total number of workgroups this dispatch schedules across all of its + /// steps (the product of each step's `dispatch_size` dimensions). Used to + /// bound how much GPU work is recorded into a single submit. + pub fn workgroup_count(&self) -> u64 { + self.steps + .iter() + .map(|step| { + let [x, y, z] = step.dispatch_size; + u64::from(x) * u64::from(y) * u64::from(z) + }) + .sum() + } + pub fn run_step<'a>(&'a self, pass: &mut ComputePass<'a>, step_index: usize) { let Some(step) = self.steps.get(step_index) else { return; diff --git a/interfaces/kalosm-chunking/Cargo.toml b/interfaces/kalosm-chunking/Cargo.toml new file mode 100644 index 000000000..1b6d253c0 --- /dev/null +++ b/interfaces/kalosm-chunking/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "kalosm-chunking" +version = "0.4.0" +edition = "2021" +description = "Text chunking primitives (sentence, paragraph, word) with language-aware sentence splitting" +license = "MIT/Apache-2.0" +repository = "https://github.com/floneum/floneum" +authors = ["Evan Almloff "] + +[dependencies] +srx = { version = "0.1.4", features = ["from_xml"] } +whatlang = "0.16.3" diff --git a/interfaces/kalosm-chunking/src/lib.rs b/interfaces/kalosm-chunking/src/lib.rs new file mode 100644 index 000000000..973a05d29 --- /dev/null +++ b/interfaces/kalosm-chunking/src/lib.rs @@ -0,0 +1,202 @@ +//! Text chunking primitives used across kalosm crates. +//! +//! This crate provides [`ChunkStrategy`] (paragraph / sentence / word sliding windows) and +//! [`SentenceChunker`] (SRX-based language-aware sentence splitting). It intentionally has +//! only two dependencies (`srx`, `whatlang`) so it can be used from wasm targets and from +//! crates that don't want the full `kalosm-language` stack. + +use std::ops::Range; + +pub mod sentence; + +pub use sentence::{DefaultSentenceChunker, SentenceChunker}; + +/// A strategy for chunking a document into smaller pieces. +/// +/// This is used to split a document into smaller pieces to generate embeddings for each piece. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum ChunkStrategy { + /// Split the document into paragraphs. + Paragraph { + /// The number of paragraphs to include in each chunk. + paragraph_count: usize, + /// The number of paragraphs to overlap between chunks. + overlap: usize, + }, + /// Split the document into sentences. + Sentence { + /// The number of sentences to include in each chunk. + sentence_count: usize, + /// The number of sentences to overlap between chunks. + overlap: usize, + }, + /// Split the document into words. + Words { + /// The number of words to include in each chunk. + word_count: usize, + /// The number of words to overlap between chunks. + overlap: usize, + }, +} + +impl ChunkStrategy { + /// Chunk a string into smaller ranges. + pub fn chunk_str(&self, string: &str) -> Vec> { + match self { + Self::Paragraph { + paragraph_count, + overlap, + } => { + let mut chunks = Vec::new(); + let mut start = 0; + let mut newline_indexes = Vec::new(); + for (i, c) in string.char_indices() { + if c == '\n' { + newline_indexes.push(i + 1); + if newline_indexes.len() >= *paragraph_count { + if !string[start..i].trim().is_empty() { + chunks.push(start..i); + } + for _ in 0..(newline_indexes.len() - *overlap) { + start = newline_indexes.remove(0); + } + } + } + } + + if !string[start..].trim().is_empty() { + chunks.push(start..string.len()); + } + + chunks + } + Self::Sentence { + sentence_count, + overlap, + } => { + let mut chunks = Vec::new(); + + let splits = SentenceChunker::default().split_sentences(string); + + for window in splits + .windows(*sentence_count) + .step_by(*sentence_count - *overlap) + { + if window.len() < *sentence_count { + break; + } + let start = window.first().unwrap().start; + let end = window.last().unwrap().end; + chunks.push(start..end); + } + + chunks + } + Self::Words { + word_count, + overlap, + } => { + let mut chunks = Vec::new(); + let mut start = 0; + let mut word_start_indexes = Vec::new(); + for (i, c) in string.char_indices() { + if c == ' ' { + word_start_indexes.push(i + 1); + if word_start_indexes.len() >= *word_count { + if !string[start..i].trim().is_empty() { + chunks.push(start..i); + } + for _ in 0..(word_start_indexes.len() - *overlap) { + start = word_start_indexes.remove(0); + } + } + } + } + + if !string[start..].trim().is_empty() { + chunks.push(start..string.len()); + } + + chunks + } + } + } +} + +impl Default for ChunkStrategy { + fn default() -> Self { + Self::Paragraph { + paragraph_count: 3, + overlap: 1, + } + } +} + +#[test] +fn test_chunking() { + let string = "The quick brown fox jumps over the lazy dog."; + let chunks = ChunkStrategy::Words { + word_count: 3, + overlap: 1, + }; + let chunks = chunks.chunk_str(string); + assert_eq!(chunks.len(), 4); + assert_eq!(string[chunks[0].clone()].trim(), "The quick brown"); + assert_eq!(string[chunks[1].clone()].trim(), "brown fox jumps"); + assert_eq!(string[chunks[2].clone()].trim(), "jumps over the"); + assert_eq!(string[chunks[3].clone()].trim(), "the lazy dog."); + + let chunks = ChunkStrategy::Words { + word_count: 3, + overlap: 2, + }; + let chunks = chunks.chunk_str(string); + assert_eq!(chunks.len(), 7); + assert_eq!(string[chunks[0].clone()].trim(), "The quick brown"); + assert_eq!(string[chunks[1].clone()].trim(), "quick brown fox"); + assert_eq!(string[chunks[2].clone()].trim(), "brown fox jumps"); + assert_eq!(string[chunks[3].clone()].trim(), "fox jumps over"); + assert_eq!(string[chunks[4].clone()].trim(), "jumps over the"); + assert_eq!(string[chunks[5].clone()].trim(), "over the lazy"); + assert_eq!(string[chunks[6].clone()].trim(), "the lazy dog."); + + let chunks = ChunkStrategy::Sentence { + sentence_count: 2, + overlap: 1, + }; + let string = "first sentence. second sentence. third sentence. fourth sentence."; + let chunks = chunks.chunk_str(string); + assert_eq!(chunks.len(), 3); + assert_eq!( + string[chunks[0].clone()].trim(), + "first sentence. second sentence." + ); + assert_eq!( + string[chunks[1].clone()].trim(), + "second sentence. third sentence." + ); + assert_eq!( + string[chunks[2].clone()].trim(), + "third sentence. fourth sentence." + ); + + let chunks = ChunkStrategy::Paragraph { + paragraph_count: 3, + overlap: 1, + }; + let string = "first paragraph\n\nsecond paragraph\n\nthird paragraph\n\nfourth paragraph"; + let chunks = chunks.chunk_str(string); + assert_eq!(chunks.len(), 3); + assert_eq!( + string[chunks[0].clone()].trim(), + "first paragraph\n\nsecond paragraph" + ); + assert_eq!( + string[chunks[1].clone()].trim(), + "second paragraph\n\nthird paragraph" + ); + assert_eq!( + string[chunks[2].clone()].trim(), + "third paragraph\n\nfourth paragraph" + ); +} diff --git a/interfaces/kalosm-language/src/search/preprocessing/sentence/assets/segment.srx b/interfaces/kalosm-chunking/src/sentence/assets/segment.srx similarity index 100% rename from interfaces/kalosm-language/src/search/preprocessing/sentence/assets/segment.srx rename to interfaces/kalosm-chunking/src/sentence/assets/segment.srx diff --git a/interfaces/kalosm-chunking/src/sentence/mod.rs b/interfaces/kalosm-chunking/src/sentence/mod.rs new file mode 100644 index 000000000..6191213ba --- /dev/null +++ b/interfaces/kalosm-chunking/src/sentence/mod.rs @@ -0,0 +1,66 @@ +use srx::SRX; +use std::cell::OnceCell; +use std::rc::Rc; +use std::str::FromStr; + +/// The default sentence chunker. Unlike [`SentenceChunker`], this is Send + Sync. +#[derive(Debug, Clone, Copy, Default)] +pub struct DefaultSentenceChunker; + +/// A sentence splitter backed by [SRX](https://www.unicode.org/uli/pas/srx/srx20.html) rules. +/// +/// Uses the [srx](https://crates.io/crates/srx) crate to parse and apply the rules. +#[derive(Debug, Clone)] +pub struct SentenceChunker { + srx: Rc, +} + +impl SentenceChunker { + /// Create a new sentence chunker from an xml rules string. + pub fn new(rules: &str) -> Self { + Self { + srx: SRX::from_str(rules) + .expect("the rules file is valid") + .into(), + } + } + + /// Create a new sentence chunker from anything that implements [`std::io::Read`] in the srx rules format. + pub fn load(reader: impl std::io::Read) -> Result { + Ok(Self { + srx: SRX::from_reader(reader)?.into(), + }) + } + + /// Split the body of a document into a list of sentence byte ranges. + pub fn split_sentences(&self, string: &str) -> Vec> { + let language = whatlang::detect_lang(string) + .map(|lang_code| lang_code.code()) + .unwrap_or("en"); + + let rules = self.srx.language_rules(language); + rules.split_ranges(string) + } +} + +impl Default for SentenceChunker { + fn default() -> Self { + // The rules are expensive to parse (~1 second), so cache them in a thread-local. + thread_local! { + static DEFAULT_RULES: OnceCell> = const { OnceCell::new() }; + } + + let rules = DEFAULT_RULES.with(|default| { + default + .get_or_init(|| { + // LanguageTool ruleset: https://github.com/languagetool-org/languagetool/blob/master/languagetool-core/src/main/resources/org/languagetool/resource/segment.srx + let rules = SRX::from_str(include_str!("./assets/segment.srx")) + .expect("the rules file is valid"); + Rc::new(rules) + }) + .clone() + }); + + Self { srx: rules } + } +} diff --git a/interfaces/kalosm-language/Cargo.toml b/interfaces/kalosm-language/Cargo.toml index fb0ea5c10..105b0dd15 100644 --- a/interfaces/kalosm-language/Cargo.toml +++ b/interfaces/kalosm-language/Cargo.toml @@ -36,11 +36,11 @@ docx-rs = { version = "0.4.7", optional = true } lopdf = { version = "0.35.0", features = ["async"], optional = true } convert_case = { version = "0.6.0", optional = true } kalosm-sample = { workspace = true } +kalosm-chunking = { workspace = true, optional = true } ego-tree = { version = "0.6.2", optional = true } image = { version = "0.24.7", optional = true } whatlang = { version = "0.16.3", optional = true } texting_robots = { version = "0.2.2", optional = true } -srx = { version = "0.1.4", features = ["from_xml"], optional = true } thiserror = { workspace = true, optional = true } anyhow = { workspace = true, optional = true } roaring = { version = "0.10.6", optional = true } @@ -74,7 +74,7 @@ chunking = [ "documents", "dep:ego-tree", "dep:scraper", - "dep:srx", + "dep:kalosm-chunking", "dep:thiserror", ] vector-db = [ diff --git a/interfaces/kalosm-language/src/search/preprocessing/chunking.rs b/interfaces/kalosm-language/src/search/preprocessing/chunking.rs index 779092abd..429830980 100644 --- a/interfaces/kalosm-language/src/search/preprocessing/chunking.rs +++ b/interfaces/kalosm-language/src/search/preprocessing/chunking.rs @@ -1,203 +1,8 @@ use kalosm_language_model::Embedder; -use std::ops::Range; -use super::{Chunker, SentenceChunker}; +use super::{ChunkStrategy, Chunker}; use crate::{prelude::Document, search::Chunk}; -/// A strategy for chunking a document into smaller pieces. -/// -/// This is used to split a document into smaller pieces to generate embeddings for each piece. -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum ChunkStrategy { - /// Split the document into paragraphs. - Paragraph { - /// The number of paragraphs to include in each chunk. - paragraph_count: usize, - /// The number of paragraphs to overlap between chunks. - overlap: usize, - }, - /// Split the document into sentences. - Sentence { - /// The number of sentences to include in each chunk. - sentence_count: usize, - /// The number of sentences to overlap between chunks. - overlap: usize, - }, - /// Split the document into words. - Words { - /// The number of words to include in each chunk. - word_count: usize, - /// The number of words to overlap between chunks. - overlap: usize, - }, -} - -impl ChunkStrategy { - /// Chunk a string into smaller ranges. - pub fn chunk_str(&self, string: &str) -> Vec> { - match self { - Self::Paragraph { - paragraph_count, - overlap, - } => { - let mut chunks = Vec::new(); - let mut start = 0; - let mut newline_indexes = Vec::new(); - for (i, c) in string.char_indices() { - if c == '\n' { - newline_indexes.push(i + 1); - if newline_indexes.len() >= *paragraph_count { - if !string[start..i].trim().is_empty() { - chunks.push(start..i); - } - for _ in 0..(newline_indexes.len() - *overlap) { - start = newline_indexes.remove(0); - } - } - } - } - - if !string[start..].trim().is_empty() { - chunks.push(start..string.len()); - } - - chunks - } - Self::Sentence { - sentence_count, - overlap, - } => { - let mut chunks = Vec::new(); - - let splits = SentenceChunker::default().split_sentences(string); - - for window in splits - .windows(*sentence_count) - .step_by(*sentence_count - *overlap) - { - if window.len() < *sentence_count { - break; - } - let start = window.first().unwrap().start; - let end = window.last().unwrap().end; - chunks.push(start..end); - } - - chunks - } - Self::Words { - word_count, - overlap, - } => { - let mut chunks = Vec::new(); - let mut start = 0; - let mut word_start_indexes = Vec::new(); - for (i, c) in string.char_indices() { - if c == ' ' { - word_start_indexes.push(i + 1); - if word_start_indexes.len() >= *word_count { - if !string[start..i].trim().is_empty() { - chunks.push(start..i); - } - for _ in 0..(word_start_indexes.len() - *overlap) { - start = word_start_indexes.remove(0); - } - } - } - } - - if !string[start..].trim().is_empty() { - chunks.push(start..string.len()); - } - - chunks - } - } - } -} - -impl Default for ChunkStrategy { - fn default() -> Self { - Self::Paragraph { - paragraph_count: 3, - overlap: 1, - } - } -} - -#[test] -fn test_chunking() { - let string = "The quick brown fox jumps over the lazy dog."; - let chunks = ChunkStrategy::Words { - word_count: 3, - overlap: 1, - }; - let chunks = chunks.chunk_str(string); - assert_eq!(chunks.len(), 4); - assert_eq!(string[chunks[0].clone()].trim(), "The quick brown"); - assert_eq!(string[chunks[1].clone()].trim(), "brown fox jumps"); - assert_eq!(string[chunks[2].clone()].trim(), "jumps over the"); - assert_eq!(string[chunks[3].clone()].trim(), "the lazy dog."); - - let chunks = ChunkStrategy::Words { - word_count: 3, - overlap: 2, - }; - let chunks = chunks.chunk_str(string); - assert_eq!(chunks.len(), 7); - assert_eq!(string[chunks[0].clone()].trim(), "The quick brown"); - assert_eq!(string[chunks[1].clone()].trim(), "quick brown fox"); - assert_eq!(string[chunks[2].clone()].trim(), "brown fox jumps"); - assert_eq!(string[chunks[3].clone()].trim(), "fox jumps over"); - assert_eq!(string[chunks[4].clone()].trim(), "jumps over the"); - assert_eq!(string[chunks[5].clone()].trim(), "over the lazy"); - assert_eq!(string[chunks[6].clone()].trim(), "the lazy dog."); - - let chunks = ChunkStrategy::Sentence { - sentence_count: 2, - overlap: 1, - }; - - let string = "first sentence. second sentence. third sentence. fourth sentence."; - - let chunks = chunks.chunk_str(string); - assert_eq!(chunks.len(), 3); - assert_eq!( - string[chunks[0].clone()].trim(), - "first sentence. second sentence." - ); - assert_eq!( - string[chunks[1].clone()].trim(), - "second sentence. third sentence." - ); - assert_eq!( - string[chunks[2].clone()].trim(), - "third sentence. fourth sentence." - ); - - let chunks = ChunkStrategy::Paragraph { - paragraph_count: 3, - overlap: 1, - }; - - let string = "first paragraph\n\nsecond paragraph\n\nthird paragraph\n\nfourth paragraph"; - - let chunks = chunks.chunk_str(string); - assert_eq!(chunks.len(), 3); - assert_eq!( - string[chunks[0].clone()].trim(), - "first paragraph\n\nsecond paragraph" - ); - assert_eq!( - string[chunks[1].clone()].trim(), - "second paragraph\n\nthird paragraph" - ); - assert_eq!( - string[chunks[2].clone()].trim(), - "third paragraph\n\nfourth paragraph" - ); -} - impl Chunker for ChunkStrategy { type Error = E; diff --git a/interfaces/kalosm-language/src/search/preprocessing/mod.rs b/interfaces/kalosm-language/src/search/preprocessing/mod.rs index 53257c947..63b63fce4 100644 --- a/interfaces/kalosm-language/src/search/preprocessing/mod.rs +++ b/interfaces/kalosm-language/src/search/preprocessing/mod.rs @@ -18,13 +18,12 @@ use crate::context::Document; use super::Chunk; +pub use kalosm_chunking::{ChunkStrategy, DefaultSentenceChunker, SentenceChunker}; mod chunking; -pub use chunking::*; mod task; pub use task::*; -mod sentence; -pub use sentence::*; mod semantic; +mod sentence; pub use semantic::*; mod html; pub use html::*; diff --git a/interfaces/kalosm-language/src/search/preprocessing/sentence/mod.rs b/interfaces/kalosm-language/src/search/preprocessing/sentence/mod.rs index fee6de195..95e0c1350 100644 --- a/interfaces/kalosm-language/src/search/preprocessing/sentence/mod.rs +++ b/interfaces/kalosm-language/src/search/preprocessing/sentence/mod.rs @@ -1,25 +1,17 @@ use crate::prelude::{Chunk, Chunker, Document, Embedder}; -use srx::SRX; -use std::cell::OnceCell; use std::future::Future; -use std::rc::Rc; -use std::str::FromStr; -/// The default sentence chunker. Unlike [`SentenceChunker`], this is Send + Sync -#[derive(Debug, Clone, Copy, Default)] -pub struct DefaultSentenceChunker; +use super::{DefaultSentenceChunker, SentenceChunker}; impl Chunker for DefaultSentenceChunker { type Error = E; - /// Chunk a document into embedded snippets. fn chunk( &self, document: &Document, embedder: &E, ) -> impl Future, Self::Error>> + Send { let default = SentenceChunker::default(); - // Split the document into sentences. We first just collect the sentences as strings and byte ranges let mut initial_chunks = Vec::new(); let body = document.body(); let ranges = default.split_sentences(document.body()); @@ -31,78 +23,14 @@ impl Chunker for DefaultSentenceChunker { } } -/// A [`Chunker`] that splits a string into sentences with a given [SRX](https://www.unicode.org/uli/pas/srx/srx20.html) rules. -/// -/// This uses the [srx](https://crates.io/crates/srx) crate to parse and apply the rules. -#[derive(Debug, Clone)] -pub struct SentenceChunker { - srx: Rc, -} - -impl SentenceChunker { - /// Create a new sentence chunker from a xml rules string - pub fn new(rules: &str) -> Self { - Self { - srx: SRX::from_str(rules) - .expect("the rules file is valid") - .into(), - } - } - - /// Create a new sentence chunker from anything that implements [`std::io::Read`](std::io::Read) in the srx rules format - pub fn load(reader: impl std::io::Read) -> Result { - Ok(Self { - srx: SRX::from_reader(reader)?.into(), - }) - } - - /// Split the body of a document into a list of ranges with sentences - pub fn split_sentences(&self, string: &str) -> Vec> { - // Try to autodetect the language of the document - let language = whatlang::detect_lang(string) - .map(|lang_code| lang_code.code()) - .unwrap_or("en"); - - // Then get the language specific rules to split the document into sentences - let rules = self.srx.language_rules(language); - - rules.split_ranges(string) - } -} - -impl Default for SentenceChunker { - fn default() -> Self { - // The rules are expensive to parse (~1 second), so we cache them in a static once cell - thread_local! { - static DEFAULT_RULES: OnceCell> = const { OnceCell::new() }; - } - - let rules = DEFAULT_RULES.with(|default| { - default - .get_or_init(|| { - // Defaults to the language tool ruleset: https://github.com/languagetool-org/languagetool/blob/master/languagetool-core/src/main/resources/org/languagetool/resource/segment.srx - let rules = SRX::from_str(include_str!("./assets/segment.srx")) - .expect("the rules file is valid"); - Rc::new(rules) - }) - .clone() - }); - - Self { srx: rules } - } -} - -/// A strategy for chunking a document into smaller pieces. impl Chunker for SentenceChunker { type Error = E; - /// Chunk a document into embedded snippets. fn chunk( &self, document: &Document, embedder: &E, ) -> impl Future, Self::Error>> + Send { - // Split the document into sentences. We first just collect the sentences as strings and byte ranges let mut initial_chunks = Vec::new(); let body = document.body(); let ranges = self.split_sentences(document.body()); @@ -119,10 +47,8 @@ async fn embed_chunk( initial_chunks: Vec, ranges: Vec>, ) -> Result, E::Error> { - // Next embed them all in one big batch let embeddings = embedder.embed_vec(initial_chunks).await?; - // Now merge the embeddings and ranges into chunks let mut chunks = Vec::new(); for (embedding, chunk) in embeddings.into_iter().zip(ranges) { let chunk = Chunk { diff --git a/models/kalosm-llama/src/gguf_tokenizer.rs b/models/kalosm-llama/src/gguf_tokenizer.rs index 95d2c037a..35d26c204 100644 --- a/models/kalosm-llama/src/gguf_tokenizer.rs +++ b/models/kalosm-llama/src/gguf_tokenizer.rs @@ -191,6 +191,17 @@ impl PreTokenizer { std::mem::swap(&mut buffers.pieces, &mut buffers.next); } } + + #[cfg(all(test, feature = "hf-tokenizer-json"))] + fn split(&self, text: &str) -> Vec { + let mut buffers = PreTokenizationBuffers::default(); + self.split_into_ranges(text, &mut buffers); + buffers + .pieces + .iter() + .map(|range| range.as_str(text).to_string()) + .collect() + } } fn split_piece(split_regex: &SplitRegex, text: &str, range: TextRange, out: &mut Vec) { @@ -833,10 +844,7 @@ mod tests { unreachable!("all bytes are covered by the byte-level alphabet") } - fn legacy_splits<'a>( - pre_tokenizer: &Sequence, - text: &'a str, - ) -> tokenizers::Result> { + fn legacy_splits(pre_tokenizer: &Sequence, text: &str) -> tokenizers::Result> { let mut pretokenized = PreTokenizedString::from(text); pre_tokenizer.pre_tokenize(&mut pretokenized)?; Ok(pretokenized diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index 104a3721b..a56f262d4 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -1,680 +1,93 @@ -use crate::raw::rope::RopeImplementation; - -use fusor::cache::AttentionMask; -use fusor::cache::KvCache; -use fusor::layers::Linear; -use fusor::layers::RmsNorm; -use fusor::QMatrix; -use fusor::Tensor; -use fusor::D; -use fusor::{CastTensor, CastTo, FloatDataType, Fusion, SimdElement}; - -pub enum FeedForwardVariant { - // Used by the Llama, Qwen, and Gemma models - Llama(Box>), - // Used by the Phi models - Phi(PhiFeedForward), -} - -impl FeedForwardVariant -where - F: CastTo + CastTensor, - f32: CastTo + CastTensor, -{ - pub(crate) fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> - where - B: Fusion<3, F>, - { - match self { - FeedForwardVariant::Llama(ffn) => ffn.forward(x), - FeedForwardVariant::Phi(ffn) => ffn.forward(x), - } - } - - pub(crate) fn forward_add_residuals( - &self, - x: &Tensor<3, F, B>, - first: &Tensor<3, f32, B1>, - second: &Tensor<3, f32, B2>, - ) -> Option> - where - B: Fusion<3, F>, - B1: Fusion<3, f32>, - B2: Fusion<3, f32>, - { - match self { - FeedForwardVariant::Llama(ffn) => ffn.forward_add_residuals(x, first, second), - FeedForwardVariant::Phi(_) => None, - } - } -} - -pub struct PhiFeedForward { - pub up: QMatrix, - pub down: QMatrix, - pub feed_forward_length: usize, -} - -impl PhiFeedForward { - pub(crate) fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> - where - F: FloatDataType + SimdElement + Default + CastTo + CastTensor, - f32: CastTo + CastTensor, - B: Fusion<3, F>, - { - // All computation happens in f32 for compatibility with SIMD ops - let x_f32 = x.cast::(); - let up_states = x_f32.q_mat_mul(&self.up); - let gate = up_states - .narrow(D::Minus1, 0, self.feed_forward_length) - .to_concrete(); - let up_states = up_states - .narrow( - D::Minus1, - self.feed_forward_length, - self.feed_forward_length, - ) - .to_concrete(); - let gate = gate.silu(); - let up_states = up_states * gate; - let result = up_states.q_mat_mul(&self.down); - result.cast() - } -} - -pub struct LlamaFeedForward { - gate: QMatrix, - gate_up: Option, - gate_bias: Option>, - down: QMatrix, - down_bias: Option>, - up: QMatrix, - up_bias: Option>, -} - -impl LlamaFeedForward { - pub(crate) fn new(gate: QMatrix, down: QMatrix, up: QMatrix) -> Self { - let gate_up = QMatrix::concat_rows(&[&gate, &up]); - Self { - gate, - gate_up, - down, - up, - gate_bias: None, - down_bias: None, - up_bias: None, - } - } - - #[cfg(feature = "vision")] - pub(crate) fn new_with_bias( - gate: QMatrix, - gate_bias: Option>, - down: QMatrix, - down_bias: Option>, - up: QMatrix, - up_bias: Option>, - ) -> Self { - let gate_up = QMatrix::concat_rows(&[&gate, &up]); - Self { - gate, - gate_up, - gate_bias, - down, - down_bias, - up, - up_bias, - } - } - - pub(crate) fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> - where - F: CastTo + CastTensor, - f32: CastTo + CastTensor, - B: Fusion<3, F>, - { - let up_result = self.activation(x); - let mut up = up_result.q_mat_mul(&self.down); - if let Some(ref bias) = self.down_bias { - let bias_f32: Tensor<1, f32> = bias.cast(); - up = up.add_(&bias_f32); - } - - // Cast back to F - up.cast() - } - - pub(crate) fn forward_add_residuals( - &self, - x: &Tensor<3, F, B>, - first: &Tensor<3, f32, B1>, - second: &Tensor<3, f32, B2>, - ) -> Option> - where - F: CastTo + CastTensor, - f32: CastTo + CastTensor, - B: Fusion<3, F>, - B1: Fusion<3, f32>, - B2: Fusion<3, f32>, - { - if self.down_bias.is_some() { - return None; - } - if x.shape()[1] > 1 { - return None; - } +//! The Llama/Qwen decoder layer is the shared [`fusor::TransformerBlock`] +//! specialized to this crate's [`RopeImplementation`]. This module wires the +//! decoder's RoPE cache into the block via [`fusor::RopeLike`], aliases the +//! block as [`LlamaAttention`], and (under the `vision` feature) provides a +//! tracing wrapper that interleaves NaN probes through the attention sublayer. - let up_result = self.activation(x); - // Residual adds authored in natural graph form: the resolver folds both - // `add`s into the qmatmul post epilogue (one dispatch on decode). - let projected = up_result.q_mat_mul(&self.down); - let with_first = (&projected + first).to_concrete(); - let up = (&with_first + second).to_concrete(); - Some(up.cast()) - } - - fn activation(&self, x: &Tensor<3, F, B>) -> Tensor<3, f32> - where - F: CastTo + CastTensor, - B: Fusion<3, F>, - { - // All computation happens in f32 for compatibility with SIMD ops - let x_f32 = x.cast::(); - - match &self.gate_up { - Some(gate_up) if self.gate_bias.is_none() && self.up_bias.is_none() => { - // SwiGLU split/gate authored in natural graph form: the resolver - // folds `silu(gate) * up` over the two narrow halves into the - // qmatmul accumulator-offset epilogue (one dispatch on decode). - let pair_len = gate_up.shape()[0] / 2; - let projected = x_f32.q_mat_mul(gate_up); - let gate = projected.narrow(D::Minus1, 0, pair_len).to_concrete(); - let up = projected - .narrow(D::Minus1, pair_len, pair_len) - .to_concrete(); - (gate.silu() * up).to_concrete() - } - Some(gate_up) => { - let gate_width = self.gate.shape()[0]; - let up_width = self.up.shape()[0]; - let gate_up_states = x_f32.q_mat_mul(gate_up); - - let mut gate_states = gate_up_states - .narrow(D::Minus1, 0, gate_width) - .to_concrete(); - if let Some(ref bias) = self.gate_bias { - let bias_f32: Tensor<1, f32> = bias.cast(); - gate_states = gate_states.add_(&bias_f32); - } - - let mut up_states = gate_up_states - .narrow(D::Minus1, gate_width, up_width) - .to_concrete(); - if let Some(ref bias) = self.up_bias { - let bias_f32: Tensor<1, f32> = bias.cast(); - up_states = up_states.add_(&bias_f32); - } - - (gate_states.silu() * up_states).to_concrete() - } - None => { - let mut w1 = x_f32.q_mat_mul(&self.gate); - if let Some(ref bias) = self.gate_bias { - let bias_f32: Tensor<1, f32> = bias.cast(); - w1 = w1.add_(&bias_f32); - } - let w1 = w1.silu(); - - let mut w3 = x_f32.q_mat_mul(&self.up); - if let Some(ref bias) = self.up_bias { - let bias_f32: Tensor<1, f32> = bias.cast(); - w3 = w3.add_(&bias_f32); - } - - (w1 * w3).to_concrete() - } - } - } -} - -pub enum AttentionVariant { - Separate(Box>), - Grouped(GroupedAttention), -} - -pub struct AttentionBias { - bias_q: Tensor<1, F>, - bias_k: Tensor<1, F>, - bias_v: Tensor<1, F>, - bias_qkv: Tensor<1, F>, -} +use crate::raw::rope::RopeImplementation; +use fusor::{CastTensor, CastTo, FloatDataType, RopeLike, SimdElement, Tensor}; -impl AttentionBias { - pub fn new(q: Tensor<1, F>, k: Tensor<1, F>, v: Tensor<1, F>) -> Self { - let bias_qkv = fusor::cat([q.clone(), k.clone(), v.clone()], 0).to_concrete(); - Self { - bias_q: q, - bias_k: k, - bias_v: v, - bias_qkv, - } - } -} +// Re-export the shared block building blocks so the rest of the crate keeps +// importing them from `attention_layer`. +pub(crate) use fusor::{ + AttentionBias, AttentionVariant, FeedForwardVariant, GroupedAttention, LlamaFeedForward, Norm, + PhiFeedForward, SeparateAttention, TransformerBlock, +}; -pub struct SeparateAttention { - pub attention_wq: QMatrix, - pub attention_qkv: Option, - pub attention_q_norm: Option>, - pub attention_wk: QMatrix, - pub attention_k_norm: Option>, - pub attention_wv: QMatrix, - pub bias: Option>, - pub interleaved_rope: bool, -} +/// The decoder layer: the shared transformer block parameterized by this +/// crate's RoPE implementation (plain Llama RoPE or Qwen-VL multi-axis RoPE). +pub(crate) type LlamaAttention = TransformerBlock>; -impl SeparateAttention +impl RopeLike for RopeImplementation where - F: CastTo + CastTensor, + F: FloatDataType + SimdElement + CastTo + CastTensor, f32: CastTo + CastTensor, { - #[allow(clippy::too_many_arguments)] - fn forward( + fn apply( &self, - num_heads: usize, - head_dim: usize, - num_key_value_heads: usize, - hidden_states: &Tensor<3, F, B>, - rope_cache: &RopeImplementation, + query: &Tensor<4, F>, + key: &Tensor<4, F>, start_pos: usize, - pos_ids: Option<&Tensor<2, F>>, - ) -> (Tensor<4, F>, Tensor<4, F>, Tensor<4, F>) - where - B: Fusion<3, F>, - { - let [b_sz, seq_len, _] = hidden_states.shape(); - - // Compute in f32 for SIMD ops compatibility - let hidden_f32 = hidden_states.cast::(); - - if let Some(attention_qkv) = &self.attention_qkv { - let query_width = num_heads * head_dim; - let key_width = num_key_value_heads * head_dim; - let value_width = key_width; - let mut qkv = hidden_f32.q_mat_mul(attention_qkv); - if let Some(bias) = &self.bias { - let bias_f32: Tensor<1, f32> = bias.bias_qkv.cast(); - qkv = qkv.add_(&bias_f32); - } - - let query_states: Tensor<4, F> = { - let query_states = qkv.narrow(D::Minus1, 0, query_width).to_concrete(); - - let query = query_states - .reshape([b_sz, seq_len, num_heads, head_dim]) - .transpose(1, 2) - .to_concrete(); - - let query: Tensor<4, F> = query.cast(); - if let Some(norm) = &self.attention_q_norm { - norm.forward_generic_4d(&query) - } else { - query - } - }; - - let key_states: Tensor<4, F> = { - let key_states = qkv.narrow(D::Minus1, query_width, key_width).to_concrete(); - - let key = key_states - .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) - .transpose(1, 2) - .to_concrete(); - - let key: Tensor<4, F> = key.cast(); - if let Some(norm) = &self.attention_k_norm { - norm.forward_generic_4d(&key) - } else { - key - } - }; - - let value_states: Tensor<4, F> = { - let value_states = qkv - .narrow(D::Minus1, query_width + key_width, value_width) - .to_concrete(); - - value_states - .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) - .transpose(1, 2) - .to_concrete() - .cast() - }; - - let (query_states, key_states) = rope_cache.forward( - &query_states, - &key_states, - start_pos, - pos_ids, - self.interleaved_rope, - ); - return (query_states, key_states, value_states); - } - - let query_states: Tensor<4, F> = { - let mut query_states = hidden_f32.q_mat_mul(&self.attention_wq); - - if let Some(bias) = &self.bias { - let bias_f32: Tensor<1, f32> = bias.bias_q.cast(); - query_states = query_states.add_(&bias_f32); - } - - let query = query_states - .reshape([b_sz, seq_len, num_heads, head_dim]) - .transpose(1, 2) - .to_concrete(); - - let query: Tensor<4, F> = query.cast(); - if let Some(norm) = &self.attention_q_norm { - norm.forward_generic_4d(&query) - } else { - query - } - }; - let key_states: Tensor<4, F> = { - let mut key_states = hidden_f32.q_mat_mul(&self.attention_wk); - - if let Some(bias) = &self.bias { - let bias_f32: Tensor<1, f32> = bias.bias_k.cast(); - key_states = key_states.add_(&bias_f32); - } - - let key = key_states - .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) - .transpose(1, 2) - .to_concrete(); - - let key: Tensor<4, F> = key.cast(); - if let Some(norm) = &self.attention_k_norm { - norm.forward_generic_4d(&key) - } else { - key - } - }; - let value_states: Tensor<4, F> = { - let mut value_states = hidden_f32.q_mat_mul(&self.attention_wv); - - if let Some(bias) = &self.bias { - let bias_f32: Tensor<1, f32> = bias.bias_v.cast(); - value_states = value_states.add_(&bias_f32); - } - - value_states - .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) - .transpose(1, 2) - .to_concrete() - .cast() - }; - - let (query_states, key_states) = rope_cache.forward( - &query_states, - &key_states, - start_pos, - pos_ids, - self.interleaved_rope, - ); - (query_states, key_states, value_states) - } -} - -pub struct GroupedAttention { - pub attention_qkv: QMatrix, - pub interleaved_rope: bool, -} - -impl GroupedAttention { - #[allow(clippy::too_many_arguments)] - fn forward( - &self, - num_heads: usize, - head_dim: usize, - num_key_value_heads: usize, - x: &Tensor<3, F, B>, - rope_cache: &RopeImplementation, - start_pos: usize, - pos_ids: Option<&Tensor<2, F>>, - ) -> (Tensor<4, F>, Tensor<4, F>, Tensor<4, F>) - where - F: FloatDataType + SimdElement + Default + CastTo + CastTensor, - f32: CastTo + CastTensor, - B: Fusion<3, F>, - { - let [b_sz, seq_len, _] = x.shape(); - // Compute in f32 for SIMD ops compatibility - let x_f32 = x.cast::(); - let qkv = x_f32.q_mat_mul(&self.attention_qkv); - - let query_pos = num_heads * head_dim; - let query_states = qkv.narrow(D::Minus1, 0, query_pos); - let key_states = qkv.narrow(D::Minus1, query_pos, num_key_value_heads * head_dim); - let value_states = qkv.narrow( - D::Minus1, - query_pos + num_key_value_heads * head_dim, - num_key_value_heads * head_dim, - ); - - let query_states: Tensor<4, F> = query_states - .reshape([b_sz, seq_len, num_heads, head_dim]) - .transpose(1, 2) - .to_concrete() - .cast(); - let key_states: Tensor<4, F> = key_states - .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) - .transpose(1, 2) - .to_concrete() - .cast(); - let value_states: Tensor<4, F> = value_states - .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) - .transpose(1, 2) - .to_concrete() - .cast(); - - let (query_states, key_states) = rope_cache.forward( - &query_states, - &key_states, - start_pos, - pos_ids, - self.interleaved_rope, - ); - - (query_states, key_states, value_states) + position_ids: Option<&Tensor<2, F>>, + interleaved: bool, + ) -> (Tensor<4, F>, Tensor<4, F>) { + self.forward(query, key, start_pos, position_ids, interleaved) } } -pub struct LlamaAttention { - pub attention_variant: AttentionVariant, - pub attention_wo: Linear, - pub attention_norm: RmsNorm<1, F>, - pub post_attention_norm: Option>, - pub feed_forward_variant: FeedForwardVariant, - pub ffn_norm: RmsNorm<1, F>, - pub post_ffn_norm: Option>, - pub n_head: usize, - pub n_kv_head: usize, - pub head_dim: usize, - pub hidden_size: usize, - pub rope_cache: RopeImplementation, - pub(crate) sliding_window_size: Option, -} - -impl LlamaAttention +/// Attention sublayer with NaN probes interleaved between each step. A free +/// function (not a block method) because it calls this crate's debug helpers, +/// which cannot live in `fusor`. Mirrors [`TransformerBlock::forward`]. +#[cfg(feature = "vision")] +#[allow(clippy::too_many_arguments)] +pub(crate) fn forward_with_trace( + block: &TransformerBlock, + hidden_states: &Tensor<3, F, B>, + attention_mask: Option<&fusor::cache::AttentionMask>, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + cache: Option<&mut fusor::cache::KvCache>, + layer_idx: usize, +) -> Tensor<3, F> where - F: CastTo + CastTensor, + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, f32: CastTo + CastTensor, + R: RopeLike, + B: fusor::Fusion<3, F>, { - pub(crate) fn forward( - &self, - hidden_states: &Tensor<3, F, B>, - attention_mask: Option<&AttentionMask>, - start_pos: usize, - pos_ids: Option<&Tensor<2, F>>, - cache: Option<&mut KvCache>, - ) -> Tensor<3, F> - where - B: Fusion<3, F>, - { - let [b_sz, q_len, _] = hidden_states.shape(); - let hidden_size = self.hidden_size; - let num_heads = self.n_head; - let head_dim = self.head_dim; - let num_key_value_heads = self.n_kv_head; - - let (query_states, key_states, value_states) = match self.attention_variant { - AttentionVariant::Separate(ref attention) => attention.forward( - num_heads, - head_dim, - num_key_value_heads, - hidden_states, - &self.rope_cache, - start_pos, - pos_ids, - ), - AttentionVariant::Grouped(ref attention) => attention.forward( - num_heads, - head_dim, - num_key_value_heads, - hidden_states, - &self.rope_cache, - start_pos, - pos_ids, - ), - }; - - // Convert to f32 for cache operations (cache uses f32 for SIMD compatibility) - let query_f32: Tensor<4, f32> = query_states.cast(); - let key_f32: Tensor<4, f32> = key_states.cast(); - let value_f32: Tensor<4, f32> = value_states.cast(); - - let (key_f32, value_f32) = match cache { - None => (key_f32, value_f32), - Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), - }; - - forward_attention_qkv_f32( - &query_f32, - &key_f32, - &value_f32, - &self.attention_wo, - attention_mask, - head_dim, - b_sz, - q_len, - hidden_size, - ) - } - - #[cfg(feature = "vision")] - pub(crate) fn forward_with_trace( - &self, - hidden_states: &Tensor<3, F, B>, - attention_mask: Option<&AttentionMask>, - start_pos: usize, - pos_ids: Option<&Tensor<2, F>>, - cache: Option<&mut KvCache>, - layer_idx: usize, - ) -> Tensor<3, F> - where - B: Fusion<3, F>, - { - let [b_sz, q_len, _] = hidden_states.shape(); - let hidden_size = self.hidden_size; - let num_heads = self.n_head; - let head_dim = self.head_dim; - let num_key_value_heads = self.n_kv_head; - - let (query_states, key_states, value_states) = match self.attention_variant { - AttentionVariant::Separate(ref attention) => attention.forward( - num_heads, - head_dim, - num_key_value_heads, - hidden_states, - &self.rope_cache, - start_pos, - pos_ids, - ), - AttentionVariant::Grouped(ref attention) => attention.forward( - num_heads, - head_dim, - num_key_value_heads, - hidden_states, - &self.rope_cache, - start_pos, - pos_ids, - ), - }; + let [b_sz, q_len, _] = hidden_states.shape(); + let hidden_size = block.hidden_size; + + let (query_states, key_states, value_states) = block.attention_variant.forward( + block.n_head, + block.head_dim, + block.n_kv_head, + hidden_states, + &block.rope_cache, + start_pos, + pos_ids, + ); - let query_f32: Tensor<4, f32> = query_states.cast(); - let key_f32: Tensor<4, f32> = key_states.cast(); - let value_f32: Tensor<4, f32> = value_states.cast(); + let query_f32: Tensor<4, f32> = query_states.cast(); + let key_f32: Tensor<4, f32> = key_states.cast(); + let value_f32: Tensor<4, f32> = value_states.cast(); - crate::raw::debug_check_nan_f32(&query_f32, layer_idx, "Q_pre_cache", start_pos); - crate::raw::debug_check_nan_f32(&key_f32, layer_idx, "K_new", start_pos); - crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_new", start_pos); + crate::raw::debug_check_nan_f32(&query_f32, layer_idx, "Q_pre_cache", start_pos); + crate::raw::debug_check_nan_f32(&key_f32, layer_idx, "K_new", start_pos); + crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_new", start_pos); - let (key_f32, value_f32) = match cache { - None => (key_f32, value_f32), - Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), - }; + let (key_f32, value_f32) = match cache { + None => (key_f32, value_f32), + Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), + }; - crate::raw::debug_check_nan_f32(&key_f32, layer_idx, "K_cache_view", start_pos); - crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_cache_view", start_pos); + crate::raw::debug_check_nan_f32(&key_f32, layer_idx, "K_cache_view", start_pos); + crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_cache_view", start_pos); - let scale = 1. / (head_dim as f64).sqrt(); - let attn_raw = query_f32.flash_attention( - &key_f32, - &value_f32, - scale as f32, - attention_mask.map(|m| { - let kind = if m.is_strict_causal() { - fusor::MaskKind::Causal - } else { - fusor::MaskKind::QKMask - }; - (m.mask(), kind) - }), - ); - crate::raw::debug_check_nan_f32(&attn_raw, layer_idx, "flash_out", start_pos); - let attn_output = attn_raw.transpose(1, 2); - let attn_output = attn_output.reshape([b_sz, q_len, hidden_size]); - let attn_output_f: Tensor<3, F> = attn_output.cast(); - let probe_in: fusor::Tensor<3, f32> = attn_output_f.clone().cast(); - crate::raw::debug_check_nan_f32(&probe_in, layer_idx, "before_wo", start_pos); - let out = self.attention_wo.forward_generic(&attn_output_f); - out - } -} - -/// Forward attention QKV computation in f32 for SIMD compatibility. -/// All intermediate computation happens in f32, with the final result cast back to F. -#[allow(clippy::too_many_arguments)] -pub(crate) fn forward_attention_qkv_f32( - query_states: &Tensor<4, f32>, - key_states: &Tensor<4, f32>, - value_states: &Tensor<4, f32>, - attention_wo: &Linear, - attention_mask: Option<&AttentionMask>, - head_dim: usize, - b_sz: usize, - q_len: usize, - hidden_size: usize, -) -> Tensor<3, F> -where - F: FloatDataType + SimdElement + Default + CastTo + CastTensor, - f32: CastTo + CastTensor, -{ - let scale = 1. / (head_dim as f64).sqrt(); - let attn_output = query_states.flash_attention( - key_states, - value_states, + let scale = 1. / (block.head_dim as f64).sqrt(); + let attn_raw = query_f32.flash_attention( + &key_f32, + &value_f32, scale as f32, attention_mask.map(|m| { let kind = if m.is_strict_causal() { @@ -685,10 +98,11 @@ where (m.mask(), kind) }), ); - - let attn_output = attn_output.transpose(1, 2); - + crate::raw::debug_check_nan_f32(&attn_raw, layer_idx, "flash_out", start_pos); + let attn_output = attn_raw.transpose(1, 2); let attn_output = attn_output.reshape([b_sz, q_len, hidden_size]); - - attention_wo.forward_generic(&attn_output.cast()) + let attn_output_f: Tensor<3, F> = attn_output.cast(); + let probe_in: fusor::Tensor<3, f32> = attn_output_f.clone().cast(); + crate::raw::debug_check_nan_f32(&probe_in, layer_idx, "before_wo", start_pos); + block.attention_wo.forward_generic(&attn_output_f) } diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index c2d8bbfe3..3850ff321 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -75,8 +75,10 @@ use attention_layer::AttentionVariant; use attention_layer::FeedForwardVariant; use attention_layer::GroupedAttention; use attention_layer::LlamaFeedForward; +use attention_layer::Norm; use attention_layer::PhiFeedForward; use attention_layer::SeparateAttention; +use attention_layer::TransformerBlock; use fusor::cache::MaskCache; use fusor::layers::Embedding; use fusor::layers::Linear; @@ -625,17 +627,17 @@ where rope.clone() }; - layers.push(LlamaAttention { + layers.push(TransformerBlock { attention_variant, attention_wo: Linear::new(attention_wo, None), - attention_norm: decode_norm(attention_norm, rms_norm_eps)?, + attention_norm: Some(Norm::Rms(decode_norm(attention_norm, rms_norm_eps)?)), post_attention_norm: post_attention_norm - .map(|norm| decode_norm(norm, rms_norm_eps)) + .map(|norm| decode_norm(norm, rms_norm_eps).map(Norm::Rms)) .transpose()?, feed_forward_variant, - ffn_norm: decode_norm(ffn_norm, rms_norm_eps)?, + ffn_norm: Norm::Rms(decode_norm(ffn_norm, rms_norm_eps)?), post_ffn_norm: ffn_post_norm - .map(|norm| decode_norm(norm, rms_norm_eps)) + .map(|norm| decode_norm(norm, rms_norm_eps).map(Norm::Rms)) .transpose()?, n_head: head_count, n_kv_head: head_count_kv, @@ -961,7 +963,11 @@ where for (i, layer) in self.layers.iter().enumerate() { let x = layer_in; let residual: Tensor<3, f32> = x.cast(); - let x = layer.attention_norm.forward_generic(&x); + let x = layer + .attention_norm + .as_ref() + .expect("decoder layers always have a pre-attention norm") + .forward(&x); if trace_layer_nan { let probe: fusor::Tensor<3, f32> = x.clone().cast(); debug_check_nan_f32(&probe, i, "post_attn_norm", index_pos); @@ -974,7 +980,8 @@ where #[cfg(feature = "vision")] { if trace_layer_nan { - layer.forward_with_trace( + attention_layer::forward_with_trace( + layer, &x, mask.as_ref(), index_pos, @@ -1008,7 +1015,7 @@ where debug_check_nan_f32(&probe, i, "attn_out", index_pos); } if let Some(post_attention_norm) = &layer.post_attention_norm { - attn = post_attention_norm.forward_generic(&attn); + attn = post_attention_norm.forward(&attn); } let attn_f32: Tensor<3, f32> = attn.cast(); @@ -1030,7 +1037,7 @@ where } let mut x = layer.feed_forward_variant.forward(&x); if let Some(post_ffn_norm) = &layer.post_ffn_norm { - x = post_ffn_norm.forward_generic(&x); + x = post_ffn_norm.forward(&x); } let x_f32: Tensor<3, f32> = x.cast(); layer_in = (x_f32 + attn_f32 + residual).cast(); diff --git a/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs b/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs index e3e721417..63d867ad8 100644 --- a/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs +++ b/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs @@ -73,7 +73,9 @@ where let trace = std::env::var_os("KALOSM_TRACE_VBLOCK").is_some(); let flush = |t: &Tensor<3, F>| { if trace { - t.as_gpu().map(|g| g.materialize_sync()); + if let Some(g) = t.as_gpu() { + g.materialize_sync(); + } } }; let xs_3d = xs.unsqueeze(0).to_concrete(); // [1, seq, dim] @@ -143,11 +145,11 @@ struct VisionAttention { } enum VisionQkv { - Fused(Linear), + Fused(Box>), Split { - q: Linear, - k: Linear, - v: Linear, + q: Box>, + k: Box>, + v: Box>, }, } @@ -173,12 +175,12 @@ where let v_b: Tensor<1, F> = vb.get("attn_v.bias", device)?.dequantize().cast(); let qkv = if let Some(qkv_weight) = QMatrix::concat_rows(&[&q_w, &k_w, &v_w]) { let qkv_bias: Tensor<1, F> = fusor::cat([q_b, k_b, v_b], 0).to_concrete(); - VisionQkv::Fused(Linear::new(qkv_weight, Some(qkv_bias))) + VisionQkv::Fused(Box::new(Linear::new(qkv_weight, Some(qkv_bias)))) } else { VisionQkv::Split { - q: Linear::new(q_w, Some(q_b)), - k: Linear::new(k_w, Some(k_b)), - v: Linear::new(v_w, Some(v_b)), + q: Box::new(Linear::new(q_w, Some(q_b))), + k: Box::new(Linear::new(k_w, Some(k_b))), + v: Box::new(Linear::new(v_w, Some(v_b))), } }; let proj = Linear::new( @@ -213,7 +215,9 @@ where VisionQkv::Fused(qkv) => { let qkv: Tensor<3, f32> = qkv.forward_generic(xs).cast(); if trace_attn { - qkv.as_gpu().map(|g| g.materialize_sync()); + if let Some(g) = qkv.as_gpu() { + g.materialize_sync(); + } tracing::info!(" qkv: {:.2?}", t_qkv.elapsed()); } let q = qkv @@ -247,7 +251,9 @@ where .reshape([seq_len, self.head_count, self.head_dim]) .to_concrete(); if trace_attn { - v.as_gpu().map(|g| g.materialize_sync()); + if let Some(g) = v.as_gpu() { + g.materialize_sync(); + } tracing::info!(" qkv: {:.2?} (split)", t_qkv.elapsed()); } (q, k, v) @@ -272,7 +278,9 @@ where let value_states = v.transpose(0, 1).unsqueeze(0).to_concrete(); let t_after_rope = Instant::now(); if trace_attn { - value_states.as_gpu().map(|g| g.materialize_sync()); + if let Some(g) = value_states.as_gpu() { + g.materialize_sync(); + } tracing::info!( " rope: {:.2?} (incl. q/k/v split + transpose)", t_qkv.elapsed() @@ -362,7 +370,9 @@ where let attn_output = attn_output.reshape([bsz, seq_len, self.embed_dim]); let output: Tensor<3, F> = self.proj.forward_generic(&attn_output.cast()); if trace_attn { - output.as_gpu().map(|g| g.materialize_sync()); + if let Some(g) = output.as_gpu() { + g.materialize_sync(); + } tracing::info!(" flash+proj: {:.2?}", t_flash.elapsed()); } diff --git a/models/kalosm-llama/src/tokenizer.rs b/models/kalosm-llama/src/tokenizer.rs index d038da8ed..82b70ecdf 100644 --- a/models/kalosm-llama/src/tokenizer.rs +++ b/models/kalosm-llama/src/tokenizer.rs @@ -9,22 +9,24 @@ pub struct LlamaTokenizer { #[derive(Clone)] enum LlamaTokenizerInner { - Gguf(GgufTokenizer), + Gguf(Box), #[cfg(feature = "hf-tokenizer-json")] - HuggingFace(tokenizers::Tokenizer), + HuggingFace(Box), } impl LlamaTokenizer { pub(crate) fn from_gguf(tokenizer: GgufTokenizer) -> Self { Self { - inner: LlamaTokenizerInner::Gguf(tokenizer), + inner: LlamaTokenizerInner::Gguf(Box::new(tokenizer)), } } #[cfg(feature = "hf-tokenizer-json")] pub(crate) fn from_hf_bytes(bytes: Vec) -> Result { Ok(Self { - inner: LlamaTokenizerInner::HuggingFace(tokenizers::Tokenizer::from_bytes(bytes)?), + inner: LlamaTokenizerInner::HuggingFace(Box::new(tokenizers::Tokenizer::from_bytes( + bytes, + )?)), }) } diff --git a/models/rbert/src/language_model.rs b/models/rbert/src/language_model.rs index 7c0142e3b..7312c1ab3 100644 --- a/models/rbert/src/language_model.rs +++ b/models/rbert/src/language_model.rs @@ -44,6 +44,24 @@ impl Bert { )) } + /// Embed a batch of sentences with a specific pooling strategy and explicit + /// control over whether the pooled vectors are L2-normalized. + pub async fn embed_batch_with_pooling_and_normalization( + &self, + inputs: Vec<&str>, + pooling: Pooling, + normalize: bool, + ) -> Result, BertError> { + let tensors = self.embed_batch_raw_with_options(inputs, pooling, normalize)?; + + let mut embeddings = Vec::with_capacity(tensors.len()); + for tensor in tensors { + embeddings.push(self.tensor_to_embedding(tensor).await?); + } + + Ok(embeddings) + } + /// Embed a sentence with a specific pooling strategy. pub async fn embed_with_pooling( &self, @@ -60,14 +78,8 @@ impl Bert { inputs: Vec<&str>, pooling: Pooling, ) -> Result, BertError> { - let tensors = self.embed_batch_raw(inputs, pooling)?; - - let mut embeddings = Vec::with_capacity(tensors.len()); - for tensor in tensors { - embeddings.push(self.tensor_to_embedding(tensor).await?); - } - - Ok(embeddings) + self.embed_batch_with_pooling_and_normalization(inputs, pooling, true) + .await } } diff --git a/models/rbert/src/lib.rs b/models/rbert/src/lib.rs index 11ef2302d..270b5a100 100644 --- a/models/rbert/src/lib.rs +++ b/models/rbert/src/lib.rs @@ -49,7 +49,8 @@ use std::sync::{Arc, RwLock}; use tokenizers::{Encoding, PaddingDirection, PaddingParams, Tokenizer}; mod language_model; -mod raw; +/// Low-level encoder implementations (standard BERT, Qwen, ModernBERT, mDeBERTa-v3). +pub mod raw; mod source; pub use crate::language_model::*; @@ -169,9 +170,9 @@ pub enum Pooling { /// An embedding model that can be either BERT or Qwen pub enum EmbeddingModel { /// A BERT-style embedding model - Bert(BertModel), + Bert(Box), /// A Qwen-style embedding model - Qwen(QwenEmbeddingModel), + Qwen(Box), } impl EmbeddingModel { @@ -346,7 +347,7 @@ impl Bert { Some("qwen3") | Some("qwen2") => { // Load Qwen embedding model let qwen_model = QwenEmbeddingModel::load(&device, &mut vb)?; - EmbeddingModel::Qwen(qwen_model) + EmbeddingModel::Qwen(Box::new(qwen_model)) } _ => { // Load BERT model (default) @@ -354,7 +355,7 @@ impl Bert { let config: Config = serde_json::from_slice(&config_bytes).map_err(BertLoadingError::LoadConfig)?; let bert_model = BertModel::load(&device, &mut vb, &config)?; - EmbeddingModel::Bert(bert_model) + EmbeddingModel::Bert(Box::new(bert_model)) } }; @@ -374,6 +375,15 @@ impl Bert { &self, sentences: Vec<&str>, pooling: Pooling, + ) -> Result>, BertError> { + self.embed_batch_raw_with_options(sentences, pooling, true) + } + + pub(crate) fn embed_batch_raw_with_options( + &self, + sentences: Vec<&str>, + pooling: Pooling, + normalize: bool, ) -> Result>, BertError> { let embedding_dim = self.model.embedding_dim(); // Approximates the quadratic attention memory cost (seq_len^2). @@ -423,7 +433,7 @@ impl Bert { } for (indices, encodings) in chunks { - let embeddings = self.embed_batch_raw_inner(encodings, pooling)?; + let embeddings = self.embed_batch_raw_inner(encodings, pooling, normalize)?; for (i, embedding) in indices.iter().zip(embeddings) { combined[*i] = Some(embedding); } @@ -435,6 +445,7 @@ impl Bert { &self, mut tokens: Vec, pooling: Pooling, + normalize: bool, ) -> Result>, BertError> { if tokens.is_empty() { return Ok(Vec::new()); @@ -484,13 +495,19 @@ impl Bert { // Cast mask u32→f32, unsqueeze to [batch, seq_len, 1] for broadcasting let mask_f32: Tensor<2, f32> = attention_mask.cast(); let mask_3d: Tensor<3, f32, _> = mask_f32.unsqueeze(2).to_concrete(); + // Broadcast mask to match embedding shape [batch, seq, hidden] + let mask_3d: Tensor<3, f32> = mask_3d.broadcast_as(shape).to_concrete(); // Zero out padding positions, sum along seq dim → [batch, hidden_dim] let masked_embeddings = (embeddings * mask_3d).to_concrete(); let summed = masked_embeddings.sum::<2>(1); // Divide by valid token count [batch, 1] — broadcasts with [batch, hidden_dim] let valid_count = mask_f32.sum_keepdim::<1>(1); let embeddings = summed.div_(&valid_count); - let embeddings = normalize_l2(&embeddings); + let embeddings = if normalize { + normalize_l2(&embeddings) + } else { + embeddings + }; Ok(embeddings .chunk(n_sentences, 0) .into_iter() @@ -509,7 +526,11 @@ impl Bert { Pooling::Last => { // With left padding, the last token is always at the final position let indexed_embeddings = embeddings.to_concrete().i((.., n_tokens - 1, ..)); - let normalized = normalize_l2(&indexed_embeddings); + let normalized = if normalize { + normalize_l2(&indexed_embeddings) + } else { + indexed_embeddings + }; Ok(normalized .chunk(n_sentences, 0) .into_iter() diff --git a/models/rbert/src/raw/mdeberta/attention.rs b/models/rbert/src/raw/mdeberta/attention.rs new file mode 100644 index 000000000..c24caad9f --- /dev/null +++ b/models/rbert/src/raw/mdeberta/attention.rs @@ -0,0 +1,343 @@ +//! mDeBERTa disentangled self-attention. +//! +//! DeBERTa uses disentangled attention with three components: +//! - Content-to-Content (c2c): Standard attention between content vectors +//! - Content-to-Position (c2p): Attention from content to relative positions +//! - Position-to-Content (p2c): Attention from relative positions to content +//! +//! The attention score is: A = c2c + c2p + p2c + +use fusor::layers::{LayerNorm, Linear}; +use fusor::{Device, Result, Tensor, VarBuilder}; + +/// Precomputed flat index tensors for the c2p and p2c gathers, valid for a +/// single `(b_sz, num_heads, seq_len)` combination. Built once per forward in +/// [`MDebertaModel::forward`] and threaded through every layer. +pub struct GatherIndices { + /// Flat source indices for c2p, shape `[b_sz * num_heads * seq_len * seq_len]`. + pub(crate) c2p: Tensor<1, u32>, + /// Flat source indices for p2c, shape `[b_sz * num_heads * seq_len * seq_len]`. + pub(crate) p2c: Tensor<1, u32>, + pub(crate) b_sz: usize, + pub(crate) num_heads: usize, + pub(crate) seq_len: usize, +} + +/// Relative position embeddings for disentangled attention. +pub struct RelativePositionEmbedding { + /// Relative position embedding table `[2 * position_buckets, hidden_size]` + /// (e.g. `[512, 768]` for DeBERTa-v3-base). Its row count is the single + /// source of truth for the bucketing geometry in `compute_gather_indices`. + embeddings: Tensor<2, f32>, + /// LayerNorm applied to embeddings (norm_rel_ebd = "layer_norm" in DeBERTa) + layer_norm: Option>, +} + +impl RelativePositionEmbedding { + /// Load with an already-loaded LayerNorm (avoids borrow issues) + pub fn load_with_norm( + device: &Device, + vb: &mut VarBuilder, + layer_norm: Option>, + ) -> Result { + let weight = vb.get("weight", device)?; + let embeddings_raw: Tensor<2, f32> = weight.dequantize(); + + // GGUF stores shape as [hidden_size, 2*max_pos] but we need [2*max_pos, hidden_size] + let [dim0, dim1] = embeddings_raw.shape(); + let embeddings = if dim0 > dim1 { + embeddings_raw.transpose(0, 1).to_concrete() + } else { + embeddings_raw + }; + + Ok(Self { + embeddings, + layer_norm, + }) + } + + /// Apply log-bucket position encoding (matches Python make_log_bucket_position). + /// positions close to 0 use linear indexing, far positions are log-bucketed. + fn make_log_bucket_position(rel_pos: i32, bucket_size: i32, max_position: i32) -> i32 { + let sign = if rel_pos > 0 { + 1 + } else if rel_pos < 0 { + -1 + } else { + 0 + }; + let mid = bucket_size / 2; + let abs_pos = if rel_pos < mid && rel_pos > -mid { + mid - 1 + } else { + rel_pos.abs() + }; + if abs_pos <= mid { + rel_pos + } else { + // log_pos = ceil(log(abs_pos/mid) / log((max-1)/mid) * (mid-1)) + mid + let ratio = (abs_pos as f32) / (mid as f32); + let max_ratio = ((max_position - 1) as f32) / (mid as f32); + let log_pos = (ratio.ln() / max_ratio.ln() * ((mid - 1) as f32)).ceil() as i32 + mid; + log_pos * sign + } + } + + /// Number of entries (`2 * position_buckets`) in the relative + /// position embedding table — this is the per-head "position dimension" of + /// the `c2p_all` / `p2c_all` attention scores before gathering. + fn num_positions(&self) -> usize { + self.embeddings.shape()[0] + } + + /// Build the flat index tensors used to gather `c2p_all` and `p2c_all` + /// along the last two dims in a single on-device `index_select`. + /// + /// The gather semantics are: + /// c2p_out[b, h, i, j] = c2p_all[b, h, i, indices[i, j]] + /// p2c_out[b, h, i, j] = p2c_all[b, h, j, indices[i, j]] + /// + /// `indices[i, j]` is the DeBERTa log-bucketed relative position index. We + /// bake the `(b, h, i|j)` outer offsets into a single 1D `u32` tensor of + /// length `b_sz * num_heads * seq_len * seq_len`, so each gather becomes: + /// source.reshape([b*h*s*p]).index_select(0, flat_idx).reshape([b,h,s,s]) + pub fn compute_gather_indices( + &self, + b_sz: usize, + num_heads: usize, + seq_len: usize, + device: &Device, + ) -> GatherIndices { + let num_pos = self.num_positions(); + // Derive the bucketing geometry from the embedding table itself rather + // than from `*.max_relative_positions` metadata. The table has + // `2 * position_buckets` rows, so `att_span` (HF `pos_ebd_size`, the + // number of buckets per side) is exactly half its height and the max + // relative distance spans the full table. Some GGUFs store + // `max_relative_positions` as the bucket count (256) and others as the + // full span (512); deriving from the table keeps every gathered index + // within `[0, num_pos)` regardless of how the converter wrote it. + let att_span = (num_pos / 2) as i32; + let bucket_size = att_span; + let max_position = num_pos as i32; + let num_positions_i = num_pos as i32; + + // Raw relative-position indices, [seq_len, seq_len]. + let mut rel = vec![0u32; seq_len * seq_len]; + for i in 0..seq_len { + for j in 0..seq_len { + let rel_pos = i as i32 - j as i32; + let bucketed = Self::make_log_bucket_position(rel_pos, bucket_size, max_position); + let idx = (bucketed + att_span).clamp(0, num_positions_i - 1) as u32; + rel[i * seq_len + j] = idx; + } + } + + let total = b_sz * num_heads * seq_len * seq_len; + let mut c2p = vec![0u32; total]; + let mut p2c = vec![0u32; total]; + for b in 0..b_sz { + for h in 0..num_heads { + let bh_offset = ((b * num_heads + h) * seq_len) * num_pos; + for i in 0..seq_len { + let row_offset_c2p = bh_offset + i * num_pos; + for j in 0..seq_len { + let rel_idx = rel[i * seq_len + j] as usize; + // c2p: key dim = i + c2p[((b * num_heads + h) * seq_len + i) * seq_len + j] = + (row_offset_c2p + rel_idx.min(num_pos - 1)) as u32; + // p2c: key dim = j (different row offset) + let row_offset_p2c = bh_offset + j * num_pos; + p2c[((b * num_heads + h) * seq_len + i) * seq_len + j] = + (row_offset_p2c + rel_idx.min(num_pos - 1)) as u32; + } + } + } + } + + GatherIndices { + c2p: Tensor::new(device, &c2p), + p2c: Tensor::new(device, &p2c), + b_sz, + num_heads, + seq_len, + } + } + + /// Get the raw relative position embedding table (normalized). + /// Returns embeddings [2*max_pos, hidden_size] + pub fn get_embeddings(&self) -> Tensor<2, f32> { + // Apply LayerNorm to embeddings (like HuggingFace get_rel_embedding) + if let Some(ref ln) = self.layer_norm { + // Add batch dimension for LayerNorm: [num_positions, hidden_size] -> [1, num_positions, hidden_size] + let emb_3d: Tensor<3, f32> = self.embeddings.unsqueeze(0).to_concrete(); + let normed = ln.forward(&emb_3d); + // Remove batch dimension + normed.squeeze(0).to_concrete() + } else { + self.embeddings.to_concrete() + } + } +} + +/// mDeBERTa disentangled self-attention with shared key attention (share_att_key=True). +pub struct MDebertaAttention { + query: Linear, + key: Linear, + value: Linear, + output: Linear, + num_heads: usize, + head_dim: usize, + scale: f32, +} + +impl MDebertaAttention { + pub fn load( + device: &Device, + vb: &mut VarBuilder, + num_heads: usize, + head_dim: usize, + ) -> Result { + let query = Linear::load(device, &mut vb.pp("query"))?; + let key = Linear::load(device, &mut vb.pp("key"))?; + let value = Linear::load(device, &mut vb.pp("value"))?; + let output = Linear::load(device, &mut vb.pp("output"))?; + + // Scale factor for disentangled attention (3 components: c2c, c2p, p2c) + // Python: scale = scaled_size_sqrt(query_layer, scale_factor) where scale_factor=3 + // This means sqrt(head_dim * 3) + let scale = 1.0 / ((head_dim as f32) * 3.0).sqrt(); + + Ok(Self { + query, + key, + value, + output, + num_heads, + head_dim, + scale, + }) + } + + /// Forward pass with disentangled attention. + /// + /// # Arguments + /// * `hidden_states` - Input [batch, seq_len, hidden_size] + /// * `rel_pos_emb` - Relative position embedding table [2*max_pos, hidden_size] + /// * `gather_idx` - Precomputed flat indices for the c2p / p2c gathers. + /// * `attention_mask` - Optional attention mask [batch, seq_len] + pub fn forward_with_indices( + &self, + hidden_states: &Tensor<3, f32>, + rel_pos_emb: &Tensor<2, f32>, + gather_idx: &GatherIndices, + attention_bias: Option<&Tensor<4, f32>>, + ) -> Tensor<3, f32> { + use super::super::utils::split_heads; + + // Compute Q, K, V projections for content and reshape to + // [batch, num_heads, seq_len, head_dim]. + let query = split_heads( + &self.query.forward(hidden_states), + self.num_heads, + self.head_dim, + ); + let key = split_heads( + &self.key.forward(hidden_states), + self.num_heads, + self.head_dim, + ); + let value = split_heads( + &self.value.forward(hidden_states), + self.num_heads, + self.head_dim, + ); + let [batch_size, _, _, _] = query.shape(); + + // === Content-to-Content attention === + let c2c_scores = query.mat_mul(&key.transpose(2, 3)); + + // === Position attention with shared Q/K projections === + // rel_pos_emb: [2*max_pos, hidden_size] -> [1, 2*max_pos, hidden_size] + let rel_emb_3d: Tensor<3, f32> = rel_pos_emb.unsqueeze(0).to_concrete(); + let pos_query = split_heads( + &self.query.forward(&rel_emb_3d), + self.num_heads, + self.head_dim, + ); + let pos_key = split_heads( + &self.key.forward(&rel_emb_3d), + self.num_heads, + self.head_dim, + ); + let num_relative_positions = pos_query.shape()[2]; + let pos_query = pos_query + .broadcast_as([ + batch_size, + self.num_heads, + num_relative_positions, + self.head_dim, + ]) + .to_concrete(); + let pos_key = pos_key + .broadcast_as([ + batch_size, + self.num_heads, + num_relative_positions, + self.head_dim, + ]) + .to_concrete(); + + // === Content-to-Position attention === + // c2p = Q @ pos_key^T -> [batch, heads, seq, 2*max_pos] + // Then gather based on relative positions + let c2p_all = query.mat_mul(&pos_key.transpose(2, 3)); + let c2p_scores = gather_by_flat_index(&c2p_all, gather_idx, &gather_idx.c2p); + + // === Position-to-Content attention === + // p2c = K @ pos_query^T -> [batch, heads, seq, 2*max_pos] + // Then gather based on transposed relative positions + let p2c_all = key.mat_mul(&pos_query.transpose(2, 3)); + let p2c_scores = gather_by_flat_index(&p2c_all, gather_idx, &gather_idx.p2c); + + // Combine: attention = (c2c + c2p + p2c) * scale + let attn_scores = c2c_scores + .add_(&c2p_scores) + .add_(&p2c_scores) + .mul_scalar(self.scale); + + // Apply attention mask (broadcast bias to [batch, 1, 1, seq_len]) + let attn_scores = if let Some(mask_bias) = attention_bias { + attn_scores.add_(mask_bias) + } else { + attn_scores + }; + + // Softmax + let attn_probs = attn_scores.softmax_last_dim::<3>(); + + // Apply attention to values and merge heads back to [batch, seq_len, hidden]. + let context = attn_probs.mat_mul(&value); + let context = super::super::utils::merge_heads(&context); + self.output.forward(&context) + } +} + +/// On-device gather used by both c2p and p2c. `flat_idx` encodes the full +/// `((b*H + h)*S + i_or_j) * P + rel[i, j]` source offset so that the gather +/// reduces to flatten → `index_select` → reshape. +fn gather_by_flat_index( + src: &Tensor<4, f32>, + shape: &GatherIndices, + flat_idx: &Tensor<1, u32>, +) -> Tensor<4, f32> { + let [b, h, s, p] = src.shape(); + debug_assert_eq!(b, shape.b_sz); + debug_assert_eq!(h, shape.num_heads); + debug_assert_eq!(s, shape.seq_len); + src.reshape([b * h * s * p]) + .index_select(0, flat_idx) + .reshape([b, h, s, s]) + .to_concrete() +} diff --git a/models/rbert/src/raw/mdeberta/config.rs b/models/rbert/src/raw/mdeberta/config.rs new file mode 100644 index 000000000..659283e3e --- /dev/null +++ b/models/rbert/src/raw/mdeberta/config.rs @@ -0,0 +1,54 @@ +//! mDeBERTa-v3 configuration from GGUF metadata. + +use fusor::{Result, VarBuilder}; + +use super::super::utils::{load_bool_or, load_f32_or, load_u32, load_u32_or}; + +/// Configuration for mDeBERTa-v3 loaded from GGUF metadata. +#[derive(Debug, Clone)] +pub struct MDebertaConfig { + /// Number of attention heads. + pub num_heads: usize, + /// Number of transformer layers. + pub num_layers: usize, + /// Dimension per attention head. + pub head_dimension: usize, + /// LayerNorm epsilon. + pub norm_eps: f32, + /// Whether the disentangled attention reuses the content Q/K projections for + /// position embeddings (`share_att_key=true`; the only supported path). + pub share_att_key: bool, +} + +impl MDebertaConfig { + /// Load configuration from GGUF metadata. + pub fn from_gguf(vb: &VarBuilder) -> Result { + let num_heads = load_u32(vb, ".attention.head_count")? as usize; + let num_layers = load_u32(vb, ".block_count")? as usize; + let hidden_size = load_u32(vb, ".embedding_length")? as usize; + + if !hidden_size.is_multiple_of(num_heads) { + return Err(fusor::Error::msg(format!( + "hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})" + ))); + } + + let head_dimension = load_u32_or(vb, ".attention.key_length", 0); + let head_dimension = if head_dimension == 0 { + hidden_size / num_heads + } else { + head_dimension as usize + }; + + let norm_eps = load_f32_or(vb, ".attention.layer_norm_epsilon", 1e-7); + let share_att_key = load_bool_or(vb, ".attention.share_att_key", true); + + Ok(Self { + num_heads, + num_layers, + head_dimension, + norm_eps, + share_att_key, + }) + } +} diff --git a/models/rbert/src/raw/mdeberta/feed_forward.rs b/models/rbert/src/raw/mdeberta/feed_forward.rs new file mode 100644 index 000000000..d9f39706c --- /dev/null +++ b/models/rbert/src/raw/mdeberta/feed_forward.rs @@ -0,0 +1,31 @@ +//! mDeBERTa Feed Forward Network. + +use fusor::layers::Linear; +use fusor::{Device, Result, Tensor, VarBuilder}; + +/// Standard GELU Feed Forward Network for mDeBERTa. +/// +/// Formula: FFN(x) = GELU(x @ W1 + b1) @ W2 + b2 +pub struct MDebertaFeedForward { + intermediate: Linear, + output: Linear, +} + +impl MDebertaFeedForward { + pub fn load(device: &Device, vb: &mut VarBuilder) -> Result { + let intermediate = Linear::load(device, &mut vb.pp("intermediate"))?; + let output = Linear::load(device, &mut vb.pp("output"))?; + + Ok(Self { + intermediate, + output, + }) + } + + pub fn forward(&self, x: &Tensor<3, f32>) -> Tensor<3, f32> { + // Intermediate: x @ W1 + b1, then GELU + let hidden = self.intermediate.forward(x).gelu(); + // Output: hidden @ W2 + b2 + self.output.forward(&hidden) + } +} diff --git a/models/rbert/src/raw/mdeberta/layer.rs b/models/rbert/src/raw/mdeberta/layer.rs new file mode 100644 index 000000000..34bd11f06 --- /dev/null +++ b/models/rbert/src/raw/mdeberta/layer.rs @@ -0,0 +1,74 @@ +//! mDeBERTa transformer layer. + +use fusor::layers::LayerNorm; +use fusor::{Device, Result, Tensor, VarBuilder}; + +use super::attention::{GatherIndices, MDebertaAttention}; +use super::feed_forward::MDebertaFeedForward; + +/// A single mDeBERTa transformer layer. +/// +/// Architecture: +/// 1. Self-attention with disentangled attention +/// 2. Add & LayerNorm +/// 3. Feed-forward network +/// 4. Add & LayerNorm +pub struct MDebertaLayer { + attention: MDebertaAttention, + attention_norm: LayerNorm<1, f32>, + feed_forward: MDebertaFeedForward, + output_norm: LayerNorm<1, f32>, +} + +impl MDebertaLayer { + pub fn load( + device: &Device, + vb: &mut VarBuilder, + num_heads: usize, + head_dim: usize, + eps: f32, + ) -> Result { + let attention = + MDebertaAttention::load(device, &mut vb.pp("attention"), num_heads, head_dim)?; + let attention_norm = LayerNorm::load(device, &mut vb.pp("attention_norm"), eps)?; + let feed_forward = MDebertaFeedForward::load(device, &mut vb.pp("ffn"))?; + let output_norm = LayerNorm::load(device, &mut vb.pp("output_norm"), eps)?; + + Ok(Self { + attention, + attention_norm, + feed_forward, + output_norm, + }) + } + + /// Forward pass through the layer with proper position attention. + /// + /// # Arguments + /// * `hidden_states` - Input [batch, seq_len, hidden_size] + /// * `rel_pos_emb` - Relative position embedding table [2*max_pos, hidden_size] + /// * `gather_idx` - Precomputed flat indices for the c2p / p2c gathers. + /// * `attention_mask` - Optional attention mask [batch, seq_len] + pub fn forward_with_rel( + &self, + hidden_states: &Tensor<3, f32>, + rel_pos_emb: &Tensor<2, f32>, + gather_idx: &GatherIndices, + attention_bias: Option<&Tensor<4, f32>>, + ) -> Tensor<3, f32> { + // Self-attention + residual + norm + let attn_output = self.attention.forward_with_indices( + hidden_states, + rel_pos_emb, + gather_idx, + attention_bias, + ); + let hidden_states = self + .attention_norm + .forward(&hidden_states.add_(&attn_output)); + + // FFN + residual + norm + let ffn_output = self.feed_forward.forward(&hidden_states); + self.output_norm.forward(&hidden_states.add_(&ffn_output)) + } +} diff --git a/models/rbert/src/raw/mdeberta/mod.rs b/models/rbert/src/raw/mdeberta/mod.rs new file mode 100644 index 000000000..18ee18d02 --- /dev/null +++ b/models/rbert/src/raw/mdeberta/mod.rs @@ -0,0 +1,12 @@ +//! mDeBERTa-v3 encoder for GLiNER-RelEx. +//! +//! mDeBERTa uses disentangled attention with relative position embeddings, +//! which differs from ModernBERT's RoPE-based attention. + +mod attention; +mod config; +mod feed_forward; +mod layer; +mod model; + +pub use model::MDebertaModel; diff --git a/models/rbert/src/raw/mdeberta/model.rs b/models/rbert/src/raw/mdeberta/model.rs new file mode 100644 index 000000000..0a63d589a --- /dev/null +++ b/models/rbert/src/raw/mdeberta/model.rs @@ -0,0 +1,139 @@ +//! mDeBERTa-v3 encoder model. + +use fusor::layers::{Embedding, LayerNorm, Linear}; +use fusor::{Device, Result, Tensor, VarBuilder}; + +use super::attention::RelativePositionEmbedding; +use super::config::MDebertaConfig; +use super::layer::MDebertaLayer; + +/// A raw synchronous mDeBERTa-v3 encoder model. This is a bidirectional +/// transformer encoder using disentangled attention with relative position +/// embeddings (DeBERTa-v3 architecture). +pub struct MDebertaModel { + /// Token embeddings + token_embeddings: Embedding, + /// Embedding LayerNorm + embedding_norm: LayerNorm<1, f32>, + /// Relative position embeddings (shared across layers, with LayerNorm) + rel_pos_embedding: RelativePositionEmbedding, + /// Transformer layers + layers: Vec, + /// Optional encoder output projection (used by the `large` variants to + /// map the encoder's 1024-dim hidden state down to the 768-dim space + /// expected by the downstream heads). Absent on base/multi variants. + output_proj: Option>, + /// Device + device: Device, + /// Number of attention heads (the only config value needed at inference). + num_heads: usize, + span: tracing::Span, +} + +impl MDebertaModel { + /// Load mDeBERTa from GGUF weights. + pub fn load(device: &Device, vb: &mut VarBuilder) -> Result { + let config = MDebertaConfig::from_gguf(vb)?; + + // The disentangled attention reuses the content Q/K projections for the + // position embeddings (the `share_att_key=true` path, which all DeBERTa-v3 + // / mDeBERTa-v3 checkpoints use). A `share_att_key=false` model ships + // separate `pos_q_proj` / `pos_key_proj` weights that we neither load nor + // apply, so fail loudly rather than silently producing wrong scores. + if !config.share_att_key { + return Err(fusor::Error::msg( + "mDeBERTa with share_att_key=false is not supported: disentangled \ + attention here reuses the content Q/K projections for position \ + embeddings and does not load separate pos_q_proj/pos_key_proj weights.", + )); + } + + let token_embeddings = Embedding::load(device, &mut vb.pp("token_embd"))?; + let embedding_norm = LayerNorm::load(device, &mut vb.pp("embd_norm"), config.norm_eps)?; + + // The `output_norm` tensor in GGUF is the LayerNorm for the relative + // position embeddings. Load it first to avoid borrow issues. + let rel_pos_norm = LayerNorm::load(device, &mut vb.pp("output_norm"), config.norm_eps).ok(); + let rel_pos_embedding = RelativePositionEmbedding::load_with_norm( + device, + &mut vb.pp("rel_pos_embd"), + rel_pos_norm, + )?; + + let mut layers = Vec::with_capacity(config.num_layers); + for i in 0..config.num_layers { + let layer = MDebertaLayer::load( + device, + &mut vb.pp(format!("blk.{i}")), + config.num_heads, + config.head_dimension, + config.norm_eps, + )?; + layers.push(layer); + } + + // Optional post-encoder projection (only present on the large variants, + // which use DeBERTa-v3-large at 1024-dim and project down to 768). + let output_proj = Linear::load(device, &mut vb.pp("output_proj")).ok(); + + Ok(Self { + token_embeddings, + embedding_norm, + rel_pos_embedding, + layers, + output_proj, + device: device.clone(), + num_heads: config.num_heads, + span: tracing::span!(tracing::Level::TRACE, "mdeberta"), + }) + } + + /// Forward pass through the model. + /// + /// # Arguments + /// * `input_ids` - Token IDs [batch, seq_len] + /// * `attention_mask` - Optional attention mask [batch, seq_len] + /// + /// # Returns + /// Hidden states [batch, seq_len, hidden_size] + pub fn forward( + &self, + input_ids: &Tensor<2, u32>, + attention_mask: Option<&Tensor<2, u32>>, + ) -> Tensor<3, f32> { + let _enter = self.span.enter(); + let [_batch_size, seq_len] = input_ids.shape(); + + let [b_sz, _] = input_ids.shape(); + let hidden_states = self.token_embeddings.forward(input_ids); + let mut hidden_states = self.embedding_norm.forward(&hidden_states); + + // Compute the flat gather indices once per forward; every layer shares them. + let gather_idx = self.rel_pos_embedding.compute_gather_indices( + b_sz, + self.num_heads, + seq_len, + &self.device, + ); + let rel_pos_emb = self.rel_pos_embedding.get_embeddings(); + let attention_bias = attention_mask.map(|mask| { + let mask_bias = super::super::utils::attention_mask_to_bias(mask); + mask_bias.unsqueeze(1).unsqueeze(1).to_concrete() + }); + + for layer in &self.layers { + hidden_states = layer.forward_with_rel( + &hidden_states, + &rel_pos_emb, + &gather_idx, + attention_bias.as_ref(), + ); + } + + if let Some(ref proj) = self.output_proj { + hidden_states = proj.forward(&hidden_states); + } + + hidden_states + } +} diff --git a/models/rbert/src/raw/mod.rs b/models/rbert/src/raw/mod.rs index 59727e07a..42083a953 100644 --- a/models/rbert/src/raw/mod.rs +++ b/models/rbert/src/raw/mod.rs @@ -16,18 +16,30 @@ mod self_output; use self_output::*; mod intermediate_layer; use intermediate_layer::*; +/// mDeBERTa-v3 raw encoder. +pub mod mdeberta; +/// ModernBERT raw encoder. +pub mod modern_bert; +/// Qwen embedding model raw implementation. pub mod qwen; +mod utils; +pub use mdeberta::MDebertaModel; +pub use modern_bert::ModernBertModel; pub use qwen::QwenEmbeddingModel; use fusor::{Device, Result, Tensor, VarBuilder}; use serde::Deserialize; use std::fmt::Debug; +/// Non-linear activation applied between the intermediate and output dense +/// layers of each [`BertLayer`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] #[serde(rename_all = "lowercase")] pub enum HiddenAct { + /// Gaussian Error Linear Unit. Gelu, + /// Rectified Linear Unit. Relu, } diff --git a/models/rbert/src/raw/modern_bert/config.rs b/models/rbert/src/raw/modern_bert/config.rs new file mode 100644 index 000000000..03118ceda --- /dev/null +++ b/models/rbert/src/raw/modern_bert/config.rs @@ -0,0 +1,86 @@ +//! ModernBERT configuration from GGUF metadata. + +use fusor::{Result, VarBuilder}; + +use super::super::utils::{load_f32_or, load_u32, load_u32_or}; + +/// Configuration for ModernBERT loaded from GGUF metadata. +#[derive(Debug, Clone)] +pub struct ModernBertConfig { + /// Number of attention heads. + pub num_heads: usize, + /// Number of key-value heads (for GQA). + pub num_kv_heads: usize, + /// Number of transformer layers. + pub num_layers: usize, + /// Dimension per attention head. + pub head_dimension: usize, + /// Maximum context length. + pub context_length: usize, + /// RoPE base frequency used by global-attention layers. + pub rope_theta: f32, + /// RoPE base frequency used by local (sliding-window) layers. Equal to + /// `rope_theta` when the model uses a single frequency. + pub local_rope_theta: f32, + /// Every Nth layer (`layer_idx % N == 0`) uses full global attention; the + /// rest use sliding-window local attention. `1` means every layer is global + /// (the default when the GGUF predates this metadata), which reproduces the + /// original all-global behaviour. + pub global_attn_every_n_layers: usize, + /// Sliding-window size for local-attention layers (ModernBERT `local_attention`, + /// e.g. 128). A query attends to keys within `local_attention / 2` positions. + /// `0` disables windowing entirely. + pub local_attention: usize, + /// LayerNorm epsilon. + pub norm_eps: f32, +} + +impl ModernBertConfig { + /// Load configuration from GGUF metadata. + pub fn from_gguf(vb: &VarBuilder) -> Result { + let num_heads = load_u32(vb, ".attention.head_count")? as usize; + let num_kv_heads = load_u32_or(vb, ".attention.head_count_kv", num_heads as u32) as usize; + let num_layers = load_u32(vb, ".block_count")? as usize; + let hidden_size = load_u32(vb, ".embedding_length")? as usize; + + if !hidden_size.is_multiple_of(num_heads) { + return Err(fusor::Error::msg(format!( + "hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})" + ))); + } + + let context_length = load_u32_or(vb, ".context_length", 8192) as usize; + let rope_theta = load_f32_or(vb, ".rope.freq_base", 10000.0); + // Local layers may use a distinct RoPE base; absent (older GGUFs) it + // falls back to the global base so behaviour is unchanged. + let local_rope_theta = load_f32_or(vb, ".rope.local_freq_base", rope_theta); + // Default 1 (every layer global) preserves the original behaviour for + // GGUFs converted before this metadata existed. + let global_attn_every_n_layers = + load_u32_or(vb, ".attention.global_attn_every_n_layers", 1).max(1) as usize; + let local_attention = load_u32_or(vb, ".attention.local_attention", 0) as usize; + let norm_eps = load_f32_or(vb, ".attention.layer_norm_rms_epsilon", 1e-6); + + // Use attention.key_length for head dimension; fall back to + // hidden_size / num_heads if not present. + let head_dimension = load_u32_or(vb, ".attention.key_length", 0); + let head_dimension = if head_dimension == 0 { + hidden_size / num_heads + } else { + head_dimension as usize + }; + + Ok(Self { + num_heads, + num_kv_heads, + num_layers, + head_dimension, + context_length, + rope_theta, + local_rope_theta, + global_attn_every_n_layers, + local_attention, + norm_eps, + }) + } +} diff --git a/models/rbert/src/raw/modern_bert/layer.rs b/models/rbert/src/raw/modern_bert/layer.rs new file mode 100644 index 000000000..f8a723a0f --- /dev/null +++ b/models/rbert/src/raw/modern_bert/layer.rs @@ -0,0 +1,194 @@ +//! ModernBERT transformer layer: a shared [`fusor::TransformerBlock`] (fused +//! QKV + RoPE + pre-norm LayerNorm + GeGLU) plus the sliding-window local +//! attention that ModernBERT alternates with global attention. + +use fusor::layers::{LayerNorm, Linear}; +use fusor::{ + AttentionVariant, Device, FeedForwardVariant, GatedActivation, GroupedAttention, + LlamaFeedForward, Norm, Result, RopeCache, Tensor, TransformerBlock, VarBuilder, +}; + +use super::super::utils::MASK_NEG_VALUE; + +/// Build an additive sliding-window bias `[seq, seq]`: `0` where the relative +/// distance `|i - j| <= window`, and a large negative value elsewhere so those +/// positions vanish after softmax. Shared across batch and heads. +fn sliding_window_bias(seq_len: usize, window: usize, device: &Device) -> Tensor<2, f32> { + let mut data = vec![0f32; seq_len * seq_len]; + for i in 0..seq_len { + for j in 0..seq_len { + if i.abs_diff(j) > window { + data[i * seq_len + j] = MASK_NEG_VALUE; + } + } + } + Tensor::new(device, &data) + .reshape([seq_len, seq_len]) + .to_concrete() +} + +/// A single ModernBERT transformer layer. +/// +/// Global layers run the shared block directly. Local layers reuse the block's +/// projections + RoPE but compute a windowed attention (a query at position `i` +/// attends only to keys within `window` positions), which the shared +/// BatchKey-masked path cannot express. +pub struct ModernBertLayer { + block: TransformerBlock, + /// Half-window for local layers; `None` selects global attention. + window: Option, + device: Device, +} + +impl ModernBertLayer { + pub fn load( + device: &Device, + vb: &mut VarBuilder, + config: &super::config::ModernBertConfig, + layer_idx: usize, + rope_cache: &RopeCache, + window: Option, + ) -> Result { + let eps = config.norm_eps; + + // Layer 0 has no attn_norm - it uses the embedding norm instead. + let attention_norm = if layer_idx > 0 { + Some(Norm::Layer(LayerNorm::load( + device, + &mut vb.pp("attn_norm"), + eps, + )?)) + } else { + None + }; + + let wqkv = vb.get("attn_qkv.weight", device)?; + let wo = vb.get("attn_output.weight", device)?; + let ffn_norm = LayerNorm::load(device, &mut vb.pp("ffn_norm"), eps)?; + + let gate_up = vb.get("ffn_gate_up.weight", device)?; + let down = vb.get("ffn_down.weight", device)?; + + let block = TransformerBlock { + attention_variant: AttentionVariant::Grouped(GroupedAttention { + attention_qkv: wqkv, + interleaved_rope: false, + }), + attention_wo: Linear::new(wo, None), + attention_norm, + post_attention_norm: None, + // ModernBERT uses a fused gate+up weight and GELU (GeGLU). + feed_forward_variant: FeedForwardVariant::Llama(Box::new( + LlamaFeedForward::from_fused_gated(gate_up, down, GatedActivation::GeLU), + )), + ffn_norm: Norm::Layer(ffn_norm), + post_ffn_norm: None, + n_head: config.num_heads, + n_kv_head: config.num_kv_heads, + head_dim: config.head_dimension, + hidden_size: config.num_heads * config.head_dimension, + rope_cache: rope_cache.clone(), + sliding_window_size: None, + }; + + Ok(Self { + block, + window, + device: device.clone(), + }) + } + + pub fn forward( + &self, + hidden_states: &Tensor<3, f32>, + mask_bias: Option<&Tensor<2, f32>>, + ) -> Tensor<3, f32> { + let [_, seq_len, _] = hidden_states.shape(); + match self.window { + // Local layer whose window actually constrains the sequence. + Some(window) if seq_len > window + 1 => { + self.forward_windowed(hidden_states, window, mask_bias) + } + // Global attention, or a window wider than the sequence: the shared + // BatchKey-masked block. + _ => self.block.forward_block(hidden_states, mask_bias), + } + } + + /// Local sliding-window attention. Reuses the shared block's pre-norm, + /// QKV+RoPE projection, output projection, and FFN, supplying its own + /// per-batch banded attention in between. + fn forward_windowed( + &self, + hidden_states: &Tensor<3, f32>, + window: usize, + pad_bias: Option<&Tensor<2, f32>>, + ) -> Tensor<3, f32> { + let block = &self.block; + + // Pre-norm (layer 0 input is already normed by the embedding norm). + let normed = match &block.attention_norm { + Some(norm) => norm.forward(hidden_states), + None => hidden_states.clone(), + }; + + // Shared projection + RoPE. + let (query_states, key_states, value_states) = block.attention_variant.forward( + block.n_head, + block.head_dim, + block.n_kv_head, + &normed, + &block.rope_cache, + 0, + None, + ); + + let [batch_size, _, seq_len, _] = query_states.shape(); + let scale = 1.0 / (block.head_dim as f32).sqrt(); + + // The band mask is a shared `[q, k]` tensor, but per-sample padding lives + // on the key axis, so the combined mask is logically `[batch, q, k]`. The + // fused flash-attention kernel only accepts a 2D mask, so we fold the band + // and each sample's padding into a per-element `QKMask` and run the batch + // as a short loop (batch is typically 1 for single-text inference). + let band = sliding_window_bias(seq_len, window, &self.device); + let mut per_batch = Vec::with_capacity(batch_size); + for b in 0..batch_size { + let q_b = query_states.narrow(0, b, 1).to_concrete(); + let k_b = key_states.narrow(0, b, 1).to_concrete(); + let v_b = value_states.narrow(0, b, 1).to_concrete(); + let mask_b = match pad_bias { + // combined[i, j] = band[i, j] + padding[b, j] + Some(pb) => { + let row = pb + .narrow(0, b, 1) + .broadcast_as([seq_len, seq_len]) + .to_concrete(); + (&band + &row).to_concrete() + } + None => band.clone(), + }; + per_batch.push(q_b.flash_attention( + &k_b, + &v_b, + scale, + Some((&mask_b, fusor::MaskKind::QKMask)), + )); + } + let attn_output = Tensor::cat(per_batch, 0); + + // Merge heads and project output. + let attn_output = attn_output + .transpose(1, 2) + .to_concrete() + .reshape([batch_size, seq_len, block.hidden_size]) + .to_concrete(); + let attn_output = block.attention_wo.forward(&attn_output); + + // Residual + pre-norm FFN + residual. + let hidden = hidden_states.add_(&attn_output); + let ffn_input = block.ffn_norm.forward(&hidden); + let ffn_output = block.feed_forward_variant.forward(&ffn_input); + hidden.add_(&ffn_output) + } +} diff --git a/models/rbert/src/raw/modern_bert/mod.rs b/models/rbert/src/raw/modern_bert/mod.rs new file mode 100644 index 000000000..2d8e7bca8 --- /dev/null +++ b/models/rbert/src/raw/modern_bert/mod.rs @@ -0,0 +1,18 @@ +//! ModernBERT/Ettin encoder implementation. +//! +//! ModernBERT uses: +//! - RoPE (Rotary Position Embeddings), with separate global/local bases +//! - Alternating global and sliding-window local attention +//! - Pre-normalization with LayerNorm +//! - GeGLU activation in FFN +//! - No token type IDs +//! +//! Each layer is the shared [`fusor::TransformerBlock`] (fused QKV + RoPE + +//! pre-norm LayerNorm + GeGLU); the sliding-window local attention lives in +//! [`layer`] alongside the block. + +mod config; +mod layer; +mod model; + +pub use model::ModernBertModel; diff --git a/models/rbert/src/raw/modern_bert/model.rs b/models/rbert/src/raw/modern_bert/model.rs new file mode 100644 index 000000000..b7faea40d --- /dev/null +++ b/models/rbert/src/raw/modern_bert/model.rs @@ -0,0 +1,111 @@ +//! ModernBERT encoder model. + +use fusor::layers::{Embedding, LayerNorm}; +use fusor::{Device, Result, RopeCache, Tensor, VarBuilder}; + +use super::super::utils::attention_mask_to_bias; +use super::config::ModernBertConfig; +use super::layer::ModernBertLayer; + +/// A raw synchronous ModernBERT (Ettin) encoder model. This is a bidirectional +/// transformer with RoPE positional embeddings, pre-normalization, and GeGLU +/// feed-forward blocks. Each layer is a shared [`fusor::TransformerBlock`]; +/// global/local attention routing is resolved per layer at load time. +pub struct ModernBertModel { + token_embeddings: Embedding, + /// Embedding norm applied after token embeddings, before first layer + embedding_norm: LayerNorm<1, f32>, + layers: Vec, + final_norm: LayerNorm<1, f32>, + span: tracing::Span, +} + +impl ModernBertModel { + /// Load ModernBERT from GGUF weights. + pub fn load(device: &Device, vb: &mut VarBuilder) -> Result { + let config = ModernBertConfig::from_gguf(vb)?; + + // Load token embeddings + let token_embeddings = Embedding::load(device, &mut vb.pp("token_embd"))?; + + // Load embedding norm (applied before first layer) + let embedding_norm = LayerNorm::load(device, &mut vb.pp("embd_norm"), config.norm_eps)?; + + // Create RoPE caches. Global and local layers may use different bases; + // when the model has a single base the two caches are identical. + let global_rope = RopeCache::new( + config.head_dimension, + config.context_length, + config.rope_theta, + device, + )?; + let local_rope = RopeCache::new( + config.head_dimension, + config.context_length, + config.local_rope_theta, + device, + )?; + let local_window = (config.local_attention > 0).then_some(config.local_attention / 2); + + // Load transformer layers, routing each to its global/local RoPE cache + // and (for local layers) sliding-window. A layer is global when + // `idx % global_attn_every_n_layers == 0`. + let mut layers = Vec::with_capacity(config.num_layers); + for i in 0..config.num_layers { + let is_global = i % config.global_attn_every_n_layers == 0; + let (rope, window) = if is_global { + (&global_rope, None) + } else { + (&local_rope, local_window) + }; + let layer = ModernBertLayer::load( + device, + &mut vb.pp(format!("blk.{i}")), + &config, + i, + rope, + window, + )?; + layers.push(layer); + } + + // Load final layer norm + let final_norm = LayerNorm::load(device, &mut vb.pp("output_norm"), config.norm_eps)?; + + Ok(Self { + token_embeddings, + embedding_norm, + layers, + final_norm, + span: tracing::span!(tracing::Level::TRACE, "modern-bert"), + }) + } + + /// Forward pass through the model. + /// + /// Returns: [batch_size, seq_len, hidden_size] + pub fn forward( + &self, + input_ids: &Tensor<2, u32>, + attention_mask: Option<&Tensor<2, u32>>, + ) -> Tensor<3, f32> { + let _enter = self.span.enter(); + // Get token embeddings + let hidden_states = self.token_embeddings.forward(input_ids); + + // Apply embedding norm (serves as pre-norm for layer 0) + let mut hidden_states = self.embedding_norm.forward(&hidden_states); + + // Convert the padding mask to an additive bias once, then reuse it for + // every layer. + let mask_bias = attention_mask.map(attention_mask_to_bias); + + // Pass through transformer layers (each routes itself global/local). + for layer in &self.layers { + hidden_states = layer.forward(&hidden_states, mask_bias.as_ref()); + } + + // Apply final layer norm + self.final_norm.forward(&hidden_states) + } +} diff --git a/models/rbert/src/raw/qwen/attention.rs b/models/rbert/src/raw/qwen/attention.rs deleted file mode 100644 index cd6fae354..000000000 --- a/models/rbert/src/raw/qwen/attention.rs +++ /dev/null @@ -1,124 +0,0 @@ -use fusor::layers::RmsNorm; -use fusor::{Device, QMatrix, Result, RopeCache, Tensor, VarBuilder}; - -/// Qwen self-attention with separate Q/K/V projections and RoPE -pub struct QwenSelfAttention { - wq: QMatrix, - wk: QMatrix, - wv: QMatrix, - wo: QMatrix, - q_norm: Option>, - k_norm: Option>, - num_heads: usize, - num_kv_heads: usize, - head_dim: usize, -} - -impl QwenSelfAttention { - pub fn load( - device: &Device, - vb: &mut VarBuilder, - num_heads: usize, - num_kv_heads: usize, - head_dim: usize, - eps: f32, - ) -> Result { - let wq = vb.get("attn_q.weight", device)?; - let wk = vb.get("attn_k.weight", device)?; - let wv = vb.get("attn_v.weight", device)?; - let wo = vb.get("attn_output.weight", device)?; - - // Optional Q/K normalization (some Qwen models have this) - let q_norm = RmsNorm::load(device, &mut vb.pp("attn_q_norm"), eps).ok(); - let k_norm = RmsNorm::load(device, &mut vb.pp("attn_k_norm"), eps).ok(); - - Ok(Self { - wq, - wk, - wv, - wo, - q_norm, - k_norm, - num_heads, - num_kv_heads, - head_dim, - }) - } - - pub fn forward( - &self, - hidden_states: &Tensor<3, f32>, - rope_cache: &RopeCache, - attention_mask: Option<&Tensor<2, u32>>, - ) -> Tensor<3, f32> { - let [b_sz, seq_len, _hidden_size] = hidden_states.shape(); - - // Compute Q, K, V projections - let mut query_states = hidden_states - .q_mat_mul(&self.wq) - .reshape([b_sz, seq_len, self.num_heads, self.head_dim]) - .transpose(1, 2) - .to_concrete(); - - let mut key_states = hidden_states - .q_mat_mul(&self.wk) - .reshape([b_sz, seq_len, self.num_kv_heads, self.head_dim]) - .transpose(1, 2) - .to_concrete(); - - let value_states = hidden_states - .q_mat_mul(&self.wv) - .reshape([b_sz, seq_len, self.num_kv_heads, self.head_dim]) - .transpose(1, 2) - .to_concrete(); - - // Apply optional Q/K normalization - if let Some(ref q_norm) = self.q_norm { - query_states = q_norm.forward_4d(&query_states); - } - if let Some(ref k_norm) = self.k_norm { - key_states = k_norm.forward_4d(&key_states); - } - - // Apply RoPE to Q and K - let (query_states, key_states) = rope_cache.forward(&query_states, &key_states, 0); - - // Scaled dot-product attention - let hidden_size = self.num_heads * self.head_dim; - let scale = 1.0 / (self.head_dim as f32).sqrt(); - - // Convert attention mask for flash attention if provided - // The mask should be [b_sz, seq_len] where 1 = valid, 0 = pad - // Flash attention expects None for no mask, or a mask tensor - // Note: We use a large negative value instead of NEG_INFINITY because - // the GPU shader path does not support inf literals. -10000 is enough to effectively - // zero out masked positions after softmax. - const MASK_NEG_VALUE: f32 = -10000.0; - let mask: Option> = attention_mask.map(|m| { - // Convert u32 mask to f32 - // 1 (valid) -> 0.0, 0 (pad) -> large negative value - let mask_f32: Tensor<2, f32> = m.cast(); - // Create ones by adding 1 to zeros - let zeros = mask_f32.zeros_like(); - let ones = (zeros + 1.0f32).to_concrete(); - // (1 - mask) * large_neg gives: valid=0, pad=large_neg - ((ones - mask_f32) * MASK_NEG_VALUE).to_concrete() - }); - - let attn_output = query_states.flash_attention( - &key_states, - &value_states, - scale, - mask.as_ref().map(|m| (m, fusor::MaskKind::BatchKeyMask)), - ); - - // Reshape and project output - let attn_output = attn_output.transpose(1, 2); - let attn_output = attn_output - .to_concrete() - .reshape([b_sz, seq_len, hidden_size]) - .to_concrete(); - - attn_output.q_mat_mul(&self.wo) - } -} diff --git a/models/rbert/src/raw/qwen/feed_forward.rs b/models/rbert/src/raw/qwen/feed_forward.rs deleted file mode 100644 index 489ccf069..000000000 --- a/models/rbert/src/raw/qwen/feed_forward.rs +++ /dev/null @@ -1,26 +0,0 @@ -use fusor::{Device, QMatrix, Result, Tensor, VarBuilder}; - -/// Qwen-style Feed Forward Network with gate/up/down projections -/// Formula: SiLU(x @ gate) * (x @ up) @ down -pub struct QwenFeedForward { - gate: QMatrix, - up: QMatrix, - down: QMatrix, -} - -impl QwenFeedForward { - pub fn load(device: &Device, vb: &mut VarBuilder) -> Result { - let gate = vb.get("ffn_gate.weight", device)?; - let up = vb.get("ffn_up.weight", device)?; - let down = vb.get("ffn_down.weight", device)?; - - Ok(Self { gate, up, down }) - } - - pub fn forward(&self, x: &Tensor<3, f32>) -> Tensor<3, f32> { - let gate = x.q_mat_mul(&self.gate); - let up = x.q_mat_mul(&self.up); - // SiLU(gate) * up, then project down - gate.silu().mul_(&up).q_mat_mul(&self.down) - } -} diff --git a/models/rbert/src/raw/qwen/layer.rs b/models/rbert/src/raw/qwen/layer.rs deleted file mode 100644 index 55fb5464e..000000000 --- a/models/rbert/src/raw/qwen/layer.rs +++ /dev/null @@ -1,58 +0,0 @@ -use fusor::layers::RmsNorm; -use fusor::{Device, Result, RopeCache, Tensor, VarBuilder}; - -use super::attention::QwenSelfAttention; -use super::feed_forward::QwenFeedForward; - -/// A single Qwen transformer layer with pre-norm architecture -pub struct QwenLayer { - attention_norm: RmsNorm<1, f32>, - attention: QwenSelfAttention, - ffn_norm: RmsNorm<1, f32>, - feed_forward: QwenFeedForward, -} - -impl QwenLayer { - pub fn load( - device: &Device, - vb: &mut VarBuilder, - num_heads: usize, - num_kv_heads: usize, - head_dim: usize, - eps: f32, - ) -> Result { - let attention_norm = RmsNorm::load(device, &mut vb.pp("attn_norm"), eps)?; - let attention = - QwenSelfAttention::load(device, vb, num_heads, num_kv_heads, head_dim, eps)?; - let ffn_norm = RmsNorm::load(device, &mut vb.pp("ffn_norm"), eps)?; - let feed_forward = QwenFeedForward::load(device, vb)?; - - Ok(Self { - attention_norm, - attention, - ffn_norm, - feed_forward, - }) - } - - pub fn forward( - &self, - hidden_states: &Tensor<3, f32>, - rope_cache: &RopeCache, - attention_mask: Option<&Tensor<2, u32>>, - ) -> Tensor<3, f32> { - // Pre-norm + attention + residual - let residual = hidden_states; - let hidden_states = self.attention_norm.forward(hidden_states); - let hidden_states = self - .attention - .forward(&hidden_states, rope_cache, attention_mask); - let hidden_states = residual.add_(&hidden_states); - - // Pre-norm + FFN + residual - let residual = &hidden_states; - let ffn_input = self.ffn_norm.forward(&hidden_states); - let ffn_output = self.feed_forward.forward(&ffn_input); - residual.add_(&ffn_output) - } -} diff --git a/models/rbert/src/raw/qwen/mod.rs b/models/rbert/src/raw/qwen/mod.rs index 17adba441..f69a4b143 100644 --- a/models/rbert/src/raw/qwen/mod.rs +++ b/models/rbert/src/raw/qwen/mod.rs @@ -1,6 +1,7 @@ -mod attention; -mod feed_forward; -mod layer; +//! Qwen embedding model (encoder-only): separate Q/K/V + optional q/k norm, +//! RoPE, pre-norm RMSNorm, SwiGLU FFN. Built on the shared +//! [`fusor::TransformerBlock`]. + mod model; pub use model::QwenEmbeddingModel; diff --git a/models/rbert/src/raw/qwen/model.rs b/models/rbert/src/raw/qwen/model.rs index f3eb7b608..25ed456c7 100644 --- a/models/rbert/src/raw/qwen/model.rs +++ b/models/rbert/src/raw/qwen/model.rs @@ -1,7 +1,10 @@ -use fusor::layers::{Embedding, RmsNorm}; -use fusor::{Device, Result, RopeCache, Tensor, VarBuilder}; +use fusor::layers::{Embedding, Linear, RmsNorm}; +use fusor::{ + AttentionVariant, Device, FeedForwardVariant, LlamaFeedForward, Norm, Result, RopeCache, + SeparateAttention, Tensor, TransformerBlock, VarBuilder, +}; -use super::layer::QwenLayer; +use super::super::utils::attention_mask_to_bias; /// Configuration for QwenEmbeddingModel loaded from GGUF metadata #[derive(Debug, Clone)] @@ -85,12 +88,66 @@ impl QwenConfig { } } +/// Build one Qwen encoder block as a shared [`TransformerBlock`]: separate +/// Q/K/V projections with optional q/k norm, RoPE, pre-norm RMSNorm, and a +/// SwiGLU feed-forward. +fn load_qwen_block( + device: &Device, + vb: &mut VarBuilder, + config: &QwenConfig, + rope_cache: &RopeCache, +) -> Result> { + let eps = config.rms_norm_eps; + + let wq = vb.get("attn_q.weight", device)?; + let wk = vb.get("attn_k.weight", device)?; + let wv = vb.get("attn_v.weight", device)?; + let wo = vb.get("attn_output.weight", device)?; + + // Optional Q/K normalization (some Qwen models have this). + let q_norm = RmsNorm::load(device, &mut vb.pp("attn_q_norm"), eps).ok(); + let k_norm = RmsNorm::load(device, &mut vb.pp("attn_k_norm"), eps).ok(); + + let attention_norm = RmsNorm::load(device, &mut vb.pp("attn_norm"), eps)?; + let ffn_norm = RmsNorm::load(device, &mut vb.pp("ffn_norm"), eps)?; + + let gate = vb.get("ffn_gate.weight", device)?; + let up = vb.get("ffn_up.weight", device)?; + let down = vb.get("ffn_down.weight", device)?; + + Ok(TransformerBlock { + attention_variant: AttentionVariant::Separate(Box::new(SeparateAttention { + attention_wq: wq, + attention_qkv: None, + attention_q_norm: q_norm, + attention_wk: wk, + attention_k_norm: k_norm, + attention_wv: wv, + bias: None, + interleaved_rope: false, + })), + attention_wo: Linear::new(wo, None), + attention_norm: Some(Norm::Rms(attention_norm)), + post_attention_norm: None, + feed_forward_variant: FeedForwardVariant::Llama(Box::new(LlamaFeedForward::new( + gate, down, up, + ))), + ffn_norm: Norm::Rms(ffn_norm), + post_ffn_norm: None, + n_head: config.num_heads, + n_kv_head: config.num_kv_heads, + head_dim: config.head_dimension, + hidden_size: config.hidden_size, + rope_cache: rope_cache.clone(), + sliding_window_size: None, + }) +} + /// Qwen embedding model (encoder-only for embeddings) pub struct QwenEmbeddingModel { token_embeddings: Embedding, - layers: Vec, + layers: Vec>, final_norm: RmsNorm<1, f32>, - rope_cache: RopeCache, pub(crate) device: Device, config: QwenConfig, } @@ -103,7 +160,7 @@ impl QwenEmbeddingModel { // Load token embeddings let token_embeddings = Embedding::load(device, &mut vb.pp("token_embd"))?; - // Create RoPE cache + // Create RoPE cache (shared across every layer) let rope_cache = RopeCache::new( config.head_dimension, config.context_length, @@ -114,14 +171,8 @@ impl QwenEmbeddingModel { // Load transformer layers let mut layers = Vec::with_capacity(config.num_layers); for i in 0..config.num_layers { - let layer = QwenLayer::load( - device, - &mut vb.pp(format!("blk.{i}")), - config.num_heads, - config.num_kv_heads, - config.head_dimension, - config.rms_norm_eps, - )?; + let layer = + load_qwen_block(device, &mut vb.pp(format!("blk.{i}")), &config, &rope_cache)?; layers.push(layer); } @@ -132,7 +183,6 @@ impl QwenEmbeddingModel { token_embeddings, layers, final_norm, - rope_cache, device: device.clone(), config, }) @@ -149,9 +199,12 @@ impl QwenEmbeddingModel { // Get token embeddings let mut hidden_states = self.token_embeddings.forward(input_ids); - // Pass through transformer layers + // Convert the padding mask to an additive bias once, then reuse it for + // every layer (each block applies it as a BatchKey mask). + let mask_bias = attention_mask.map(attention_mask_to_bias); + for layer in &self.layers { - hidden_states = layer.forward(&hidden_states, &self.rope_cache, attention_mask); + hidden_states = layer.forward_block(&hidden_states, mask_bias.as_ref()); } // Apply final layer norm diff --git a/models/rbert/src/raw/self_attention.rs b/models/rbert/src/raw/self_attention.rs index 5bf1aef73..8319501b7 100644 --- a/models/rbert/src/raw/self_attention.rs +++ b/models/rbert/src/raw/self_attention.rs @@ -34,15 +34,7 @@ impl BertSelfAttention { } pub(crate) fn transpose_for_scores(&self, xs: &Tensor<3, f32>) -> Tensor<4, f32> { - let shape = xs.shape(); - let new_x_shape = [ - shape[0], - shape[1], - self.num_attention_heads, - self.attention_head_size, - ]; - - xs.reshape(new_x_shape).transpose(1, 2).to_concrete() + super::utils::split_heads(xs, self.num_attention_heads, self.attention_head_size) } pub(crate) fn forward( @@ -60,31 +52,16 @@ impl BertSelfAttention { let key_layer = self.transpose_for_scores(&key_layer); let value_layer = self.transpose_for_scores(&value_layer); - let context_layer = if let Some(attention_mask) = attention_mask { - let attention_scores = query_layer.mat_mul(&key_layer.t()); - let mut attention_scores = attention_scores.mul_scalar(scale); - - // If there is an attention mask, filter the attention scores by that mask - // The attention mask is a tensor of shape (bsize, seq_len) - // the attention scores are a tensor of shape (bsize, _, seq_len, seq_len) - // We expand the attention mask to (bsize, 1, 1, seq_len) - let mask = attention_mask - .unsqueeze::<3>(1) - .unsqueeze::<4>(2) - .to_concrete(); - let shape = attention_scores.shape(); - let mask: Tensor<4, f32> = mask.broadcast_as::<4>(shape).to_concrete().cast(); - // We use a value slightly larger that the true f32 min value to avoid NaN - const FALSE_MIN: f32 = -3.4028235e34f32; - let device = attention_scores.device(); - let on_false = Tensor::splat(&device, FALSE_MIN, shape); - attention_scores = mask.where_cond(&attention_scores, &on_false); + let mask = attention_mask.map(super::utils::attention_mask_to_bias); + let context_layer = { let _enter_sm = self.span_softmax.enter(); - let attention_probs = attention_scores.softmax_last_dim::<3>(); - attention_probs.mat_mul(&value_layer) - } else { - query_layer.flash_attention(&key_layer, &value_layer, scale, None) + query_layer.flash_attention( + &key_layer, + &value_layer, + scale, + mask.as_ref().map(|m| (m, fusor::MaskKind::BatchKeyMask)), + ) }; let context_layer = context_layer.transpose(1, 2).to_concrete(); context_layer.flatten_last_n::<1, _>() diff --git a/models/rbert/src/raw/utils.rs b/models/rbert/src/raw/utils.rs new file mode 100644 index 000000000..16b9b70a4 --- /dev/null +++ b/models/rbert/src/raw/utils.rs @@ -0,0 +1,73 @@ +//! Shared helpers used by the raw encoder implementations. + +use fusor::{Result, Tensor, VarBuilder}; + +/// Reshape `[batch, seq_len, num_heads * head_dim]` into +/// `[batch, num_heads, seq_len, head_dim]` for multi-head attention. +pub(crate) fn split_heads( + tensor: &Tensor<3, f32>, + num_heads: usize, + head_dim: usize, +) -> Tensor<4, f32> { + let [batch, seq_len, _] = tensor.shape(); + tensor + .reshape([batch, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete() +} + +/// Inverse of [`split_heads`]: collapse +/// `[batch, num_heads, seq_len, head_dim]` back to +/// `[batch, seq_len, num_heads * head_dim]`. +pub(crate) fn merge_heads(tensor: &Tensor<4, f32>) -> Tensor<3, f32> { + let [batch, num_heads, seq_len, head_dim] = tensor.shape(); + tensor + .transpose(1, 2) + .to_concrete() + .reshape([batch, seq_len, num_heads * head_dim]) + .to_concrete() +} + +/// Large negative bias applied to masked-out positions so that, after +/// softmax, their probability collapses to ~0. +pub(crate) const MASK_NEG_VALUE: f32 = -10000.0; + +/// Convert a `[batch, seq]` boolean attention mask (1 = attend, 0 = pad) +/// into an additive bias tensor of the same shape (0 for real tokens, +/// [`MASK_NEG_VALUE`] for padding). +pub(crate) fn attention_mask_to_bias(mask: &Tensor<2, u32>) -> Tensor<2, f32> { + let mask_f32: Tensor<2, f32> = mask.cast(); + let zeros = mask_f32.zeros_like(); + let ones = (zeros + 1.0f32).to_concrete(); + ((ones - mask_f32) * MASK_NEG_VALUE).to_concrete() +} + +/// Read a required `u32` GGUF metadata value. The `.`-prefix on keys is +/// interpreted by fusor as a suffix match, so callers pass architecture-agnostic +/// keys like `.attention.head_count`. +pub(crate) fn load_u32(vb: &VarBuilder, key: &str) -> Result { + vb.get_metadata(key) + .and_then(|v| v.to_u32().ok()) + .ok_or_else(|| fusor::Error::msg(format!("Missing required GGUF metadata: {key}"))) +} + +/// Read an optional `u32` GGUF metadata value, falling back to `default`. +pub(crate) fn load_u32_or(vb: &VarBuilder, key: &str, default: u32) -> u32 { + vb.get_metadata(key) + .and_then(|v| v.to_u32().ok()) + .unwrap_or(default) +} + +/// Read an optional `f32` GGUF metadata value, falling back to `default`. +pub(crate) fn load_f32_or(vb: &VarBuilder, key: &str, default: f32) -> f32 { + vb.get_metadata(key) + .and_then(|v| v.to_f32().ok()) + .unwrap_or(default) +} + +/// Read an optional `bool` GGUF metadata value, falling back to `default`. +pub(crate) fn load_bool_or(vb: &VarBuilder, key: &str, default: bool) -> bool { + vb.get_metadata(key) + .and_then(|v| v.to_bool().ok()) + .unwrap_or(default) +} diff --git a/models/rgliner/.gitignore b/models/rgliner/.gitignore new file mode 100644 index 000000000..7a60b85e1 --- /dev/null +++ b/models/rgliner/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/models/rgliner/Cargo.toml b/models/rgliner/Cargo.toml new file mode 100644 index 000000000..8ffacfdb2 --- /dev/null +++ b/models/rgliner/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "rgliner" +version = "0.4.0" +edition = "2021" +description = "GLiNER bi-encoder Named Entity Recognition for Rust" +license = "MIT/Apache-2.0" +repository = "https://github.com/floneum/floneum" +authors = ["Evan Almloff "] +keywords = ["ai", "ner", "nlp", "gliner", "transformers"] + +[dependencies] +fusor = { workspace = true, default-features = false } +tokenizers = { workspace = true, features = ["fancy-regex"] } +thiserror.workspace = true +rbert = { workspace = true } + +serde_json = "1.0.106" +serde = { version = "1", features = ["derive"] } + +kalosm-common = { workspace = true } +kalosm-model-types.workspace = true +kalosm-language-model.workspace = true + +[dev-dependencies] +anyhow.workspace = true +tokio = { version = "1", features = ["full"] } +clap = { version = "4", features = ["derive"] } + +[features] +default = ["cpu", "gpu"] +cpu = ["fusor/cpu", "rbert/cpu"] +gpu = ["fusor/gpu", "rbert/gpu"] diff --git a/models/rgliner/examples/basic.rs b/models/rgliner/examples/basic.rs new file mode 100644 index 000000000..873e8ba51 --- /dev/null +++ b/models/rgliner/examples/basic.rs @@ -0,0 +1,47 @@ +use rgliner::*; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + // Use local GGUF files if GLINER_MODEL env var is set, otherwise try official HuggingFace edge variant + let source = if let Ok(model_path) = std::env::var("GLINER_MODEL") { + // Derive label encoder path from model path + let label_encoder_path = model_path.replace(".gguf", "-label-encoder.gguf"); + println!("Loading GLiNER model from: {}", model_path); + println!("Loading label encoder from: {}", label_encoder_path); + GlinerSource::local(model_path, label_encoder_path) + } else { + // Use official HuggingFace GGUF + println!("Loading GLiNER model from HuggingFace (edge variant)..."); + GlinerSource::edge() + }; + + let mut gliner = Gliner::builder() + .with_source(source) + .with_threshold(0.01) + .build() + .await?; + + println!("Model loaded!"); + + let labels = ["person", "award", "date", "competitions", "teams"]; + + // The Ronaldo paragraph from the v2.0 model card (short strings score poorly on v2.0). + let texts = [ + "Cristiano Ronaldo dos Santos Aveiro (Portuguese pronunciation: [kɾiʃˈtjɐnu ʁɔˈnaldu]; born 5 February 1985) is a Portuguese professional footballer who plays as a forward for and captains both Saudi Pro League club Al Nassr and the Portugal national team.", + ]; + + for text in texts { + println!("\n--- Testing: {} ---", text); + let entities = gliner.extract(text, &labels).await?; + + println!("Found {} entities:", entities.len()); + for entity in entities { + println!( + " {}: '{}' (score: {:.2})", + entity.label, entity.text, entity.score + ); + } + } + + Ok(()) +} diff --git a/models/rgliner/examples/relex.rs b/models/rgliner/examples/relex.rs new file mode 100644 index 000000000..9ba55731e --- /dev/null +++ b/models/rgliner/examples/relex.rs @@ -0,0 +1,99 @@ +//! Example of using GlinerRelEx for joint NER and relation extraction. +//! +//! Examples: +//! ``` +//! cargo run --example relex -p rgliner --release -- \ +//! --text "Apple was founded by Steve Jobs in California." \ +//! --entity-labels person,organization,location \ +//! --relation-labels "founded by,located in" +//! ``` +//! +//! The example uses the default GLiNER-RelEx source, which downloads the +//! default GGUF model from Hugging Face. + +use clap::Parser; +use rgliner::relex::GlinerRelEx; + +#[derive(Parser, Debug)] +#[command( + about = "GLiNER-RelEx joint NER and relation extraction", + long_about = None, +)] +struct Args { + /// Input text to analyze. + #[arg(short, long)] + text: String, + + /// Entity labels to detect (comma-separated). + #[arg(short = 'e', long, value_delimiter = ',', required = true)] + entity_labels: Vec, + + /// Relation labels to detect (comma-separated). If empty, only entities are returned. + #[arg(short = 'r', long, value_delimiter = ',', default_value = "")] + relation_labels: Vec, + + /// Minimum confidence for entity detection. + #[arg(long, default_value_t = 0.5)] + entity_threshold: f32, + + /// Minimum confidence for relation classification. + #[arg(long, default_value_t = 0.5)] + relation_threshold: f32, +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let args = Args::parse(); + + // `value_delimiter = ','` with `default_value = ""` produces a single empty + // string when the user passes nothing - filter it out. + let entity_labels: Vec<&str> = args + .entity_labels + .iter() + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .collect(); + let relation_labels: Vec<&str> = args + .relation_labels + .iter() + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .collect(); + + if entity_labels.is_empty() { + anyhow::bail!("--entity-labels must contain at least one non-empty label"); + } + + println!("Loading default GLiNER-RelEx model..."); + println!("Text: {}", args.text); + println!("Entity labels: {:?}", entity_labels); + println!("Relation labels: {:?}", relation_labels); + + let relex = GlinerRelEx::builder() + .with_entity_threshold(args.entity_threshold) + .with_relation_threshold(args.relation_threshold) + .build() + .await?; + + let (entities, relations) = relex + .extract(&args.text, &entity_labels, &relation_labels) + .await?; + + println!("\nEntities found ({}):", entities.len()); + for entity in &entities { + println!( + " {} [{}] (score: {:.3})", + entity.text, entity.label, entity.score + ); + } + + println!("\nRelations found ({}):", relations.len()); + for relation in &relations { + println!( + " {} --[{}]--> {} (score: {:.3})", + relation.head.text, relation.relation, relation.tail.text, relation.score + ); + } + + Ok(()) +} diff --git a/models/rgliner/scripts/convert_relex_to_gguf.py b/models/rgliner/scripts/convert_relex_to_gguf.py new file mode 100644 index 000000000..37b2e6b4f --- /dev/null +++ b/models/rgliner/scripts/convert_relex_to_gguf.py @@ -0,0 +1,378 @@ +#!/usr/bin/env python3 +""" +Convert GLiNER-RelEx PyTorch models to GGUF format. + +Usage: + python convert_relex_to_gguf.py --model knowledgator/gliner-relex-multi-v1.0 --output gliner-relex-multi-v1.0.gguf + +This script converts the GLiNER-RelEx model including: +1. mDeBERTa-v3 encoder +2. Span representation layer +3. Relations representation layer (adjacency scoring) +4. Pair projector +5. Entity/relation label FFNs +""" + +import argparse +import json +import os +import shutil +import struct +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch +from huggingface_hub import snapshot_download + +# GGUF/GGML constants, quant tables and the GGUFWriter live in the shared module. +from gguf_common import * # noqa: F401,F403 + + +def load_pytorch_model(model_id: str, cache_dir: str = None) -> Tuple[Dict[str, torch.Tensor], Dict, str, str]: + """Load PyTorch model from HuggingFace. + + Returns: + state_dict, config (parsed), gliner_config_json (raw text), tokenizer_json (raw text) + """ + print(f"Downloading model: {model_id}") + + model_dir = snapshot_download( + model_id, + cache_dir=cache_dir, + allow_patterns=["*.bin", "*.json", "*.safetensors"] + ) + + config_path = os.path.join(model_dir, "gliner_config.json") + with open(config_path, 'r') as f: + gliner_config_json = f.read() + config = json.loads(gliner_config_json) + + tokenizer_path = os.path.join(model_dir, "tokenizer.json") + with open(tokenizer_path, 'r') as f: + tokenizer_json = f.read() + + weights_path = os.path.join(model_dir, "pytorch_model.bin") + if os.path.exists(weights_path): + print(f"Loading weights from: {weights_path}") + state_dict = torch.load(weights_path, map_location='cpu', weights_only=True) + else: + from safetensors.torch import load_file + weights_path = os.path.join(model_dir, "model.safetensors") + print(f"Loading weights from: {weights_path}") + state_dict = load_file(weights_path) + + return state_dict, config, gliner_config_json, tokenizer_json + + +def map_relex_weight_name(pytorch_name: str) -> str: + """Map PyTorch weight names to GGUF conventions for GLiNER-RelEx.""" + name = pytorch_name + + # ===== Encoder output projection (large variants: 1024 -> 768) ===== + name = name.replace("token_rep_layer.projection.weight", "text.output_proj.weight") + name = name.replace("token_rep_layer.projection.bias", "text.output_proj.bias") + + # ===== mDeBERTa Encoder (token_rep_layer.bert_layer.model.*) ===== + name = name.replace("token_rep_layer.bert_layer.model.embeddings.word_embeddings.weight", "text.token_embd.weight") + name = name.replace("token_rep_layer.bert_layer.model.embeddings.LayerNorm.weight", "text.embd_norm.weight") + name = name.replace("token_rep_layer.bert_layer.model.embeddings.LayerNorm.bias", "text.embd_norm.bias") + + # Relative position embeddings + name = name.replace("token_rep_layer.bert_layer.model.encoder.rel_embeddings.weight", "text.rel_pos_embd.weight") + name = name.replace("token_rep_layer.bert_layer.model.encoder.LayerNorm.weight", "text.output_norm.weight") + name = name.replace("token_rep_layer.bert_layer.model.encoder.LayerNorm.bias", "text.output_norm.bias") + + # DeBERTa layers + name = name.replace("token_rep_layer.bert_layer.model.encoder.layer.", "text.blk.") + + # DeBERTa attention + name = name.replace(".attention.self.query_proj.", ".attention.query.") + name = name.replace(".attention.self.key_proj.", ".attention.key.") + name = name.replace(".attention.self.value_proj.", ".attention.value.") + name = name.replace(".attention.self.pos_proj.", ".attention.pos_proj.") + name = name.replace(".attention.self.pos_q_proj.", ".attention.pos_q_proj.") + name = name.replace(".attention.output.dense.", ".attention.output.") + name = name.replace(".attention.output.LayerNorm.", ".attention_norm.") + + # DeBERTa FFN + name = name.replace(".intermediate.dense.", ".ffn.intermediate.") + name = name.replace(".output.dense.", ".ffn.output.") + name = name.replace(".output.LayerNorm.", ".output_norm.") + + # ===== Span Representation Layer ===== + name = name.replace("span_rep_layer.span_rep_layer.project_start.0.", "span.start_fc1.") + name = name.replace("span_rep_layer.span_rep_layer.project_start.3.", "span.start_fc2.") + name = name.replace("span_rep_layer.span_rep_layer.project_end.0.", "span.end_fc1.") + name = name.replace("span_rep_layer.span_rep_layer.project_end.3.", "span.end_fc2.") + name = name.replace("span_rep_layer.span_rep_layer.out_project.0.", "span.out_fc1.") + name = name.replace("span_rep_layer.span_rep_layer.out_project.3.", "span.out_fc2.") + + # ===== BiLSTM (rnn.lstm.*) ===== + name = name.replace("rnn.lstm.", "rnn.") + + # ===== Scorer (scorer.*) ===== + # Keep scorer.* as-is: scorer.proj_token, scorer.proj_label, scorer.out_mlp + + # ===== Pair Representation Layer (pair_rep_layer.*) ===== + name = name.replace("pair_rep_layer.0.", "pair_proj.0.") + name = name.replace("pair_rep_layer.3.", "pair_proj.3.") + + # ===== Prompt Representation Layer (prompt_rep_layer.*) ===== + # Keep prompt_rep_layer.* as-is + + return name + + +def _llama_quantize(f32_path: str, output_path: str, quant_type: str) -> None: + """Run `llama-quantize` to convert an f32 GGUF to a k-quant type.""" + binary = shutil.which("llama-quantize") + if binary is None: + raise RuntimeError( + "`llama-quantize` not found in PATH. Install llama.cpp to use k-quant types:\n" + " brew install llama.cpp # macOS\n" + " or build from https://github.com/ggml-org/llama.cpp" + ) + cmd = [binary] + # Keep `scorer.out_mlp.3.weight` at F32 - fusor's quantised matmul produces + # NaN on all-but-first rows when the output dim is tiny (shape is [3, 3072]). + cmd += ["--tensor-type", "scorer.out_mlp.3.weight=f32"] + cmd += [f32_path, output_path, quant_type.upper()] + print(f"\n$ {' '.join(cmd)}") + result = subprocess.run(cmd) + if result.returncode != 0: + raise RuntimeError( + f"llama-quantize failed (exit {result.returncode}) for {quant_type!r}. " + f"See output above. The f32 GGUF was kept at {f32_path} for inspection." + ) + + +def convert_relex_to_gguf( + model_id: str, + output_path: str, + quantize: str = "f32", + cache_dir: str = None +): + """Convert GLiNER-RelEx model to GGUF format.""" + + state_dict, config, gliner_config_json, tokenizer_json = load_pytorch_model(model_id, cache_dir) + + print("\nModel weights:") + for name, tensor in state_dict.items(): + print(f" {name}: {tensor.shape} {tensor.dtype}") + + quantize = quantize.lower() + if quantize not in QUANT_TYPES: + raise ValueError( + f"Unsupported quantization: {quantize!r}. Supported: {', '.join(sorted(QUANT_TYPES))}" + ) + post_quant = None + if quantize in LLAMA_QUANT_TYPES: + post_quant = quantize + quantize = "f32" + default_ggml_type = _ggml_type_for(quantize) + + # For k-quant targets, write f32 first to a temp path then shell out. + final_output_path = output_path + if post_quant is not None: + tmp = tempfile.NamedTemporaryFile( + prefix=os.path.basename(output_path).rsplit(".", 1)[0] + ".f32.", + suffix=".gguf", + delete=False, + dir=os.path.dirname(os.path.abspath(output_path)) or None, + ) + output_path = tmp.name + tmp.close() + + writer = GGUFWriter(output_path) + + # Masquerade as `bert` when we'll feed this to llama-quantize so its loader + # accepts the custom arch. rgliner doesn't read general.architecture. + arch = "bert" if post_quant is not None else "gliner-relex" + writer.add_metadata("general.architecture", arch) + writer.add_metadata("general.name", model_id.split("/")[-1]) + writer.add_metadata("general.quantization_version", 2) + + # GLiNER-RelEx specific metadata + writer.add_metadata("gliner.max_width", config.get("max_width", 12)) + writer.add_metadata("gliner.span_mode", config.get("span_mode", "markerV0")) + writer.add_metadata("gliner.subtoken_pooling", config.get("subtoken_pooling", "first")) + + # mDeBERTa config + encoder_config = config.get("encoder_config", {}) + hidden_size = encoder_config.get("hidden_size", 768) + num_heads = encoder_config.get("num_attention_heads", 12) + num_layers = encoder_config.get("num_hidden_layers", 12) + intermediate_size = encoder_config.get("intermediate_size", 3072) + vocab_size = encoder_config.get("vocab_size", 250105) + context_length = encoder_config.get("max_position_embeddings", 512) + max_relative_positions = encoder_config.get("max_relative_positions", 512) + + # Handle -1 which means "use full context" + if max_relative_positions <= 0: + # Derive from rel_embeddings shape if available, otherwise use context_length // 2 + rel_emb_key = "token_rep_layer.bert_layer.model.encoder.rel_embeddings.weight" + if rel_emb_key in state_dict: + num_positions = state_dict[rel_emb_key].shape[0] + max_relative_positions = num_positions // 2 + print(f"Derived max_relative_positions from rel_embeddings: {max_relative_positions}") + else: + max_relative_positions = context_length // 2 + print(f"Using default max_relative_positions: {max_relative_positions}") + + writer.add_metadata("gliner.attention.head_count", num_heads) + writer.add_metadata("gliner.block_count", num_layers) + writer.add_metadata("gliner.embedding_length", hidden_size) + writer.add_metadata("gliner.feed_forward_length", intermediate_size) + writer.add_metadata("gliner.context_length", context_length) + writer.add_metadata("gliner.attention.max_relative_positions", max_relative_positions) + writer.add_metadata("gliner.attention.layer_norm_epsilon", 1e-7) + writer.add_metadata("gliner.vocab_size", vocab_size) + + # RelEx-specific metadata + writer.add_metadata("gliner.relex.ent_token_id", 250102) # <> + writer.add_metadata("gliner.relex.rel_token_id", 250104) # <> + + # Embed tokenizer.json and gliner_config.json as string metadata so the + # GGUF is self-contained (no separate files needed at inference time). + writer.add_metadata("gliner.tokenizer_json", tokenizer_json) + writer.add_metadata("gliner.config_json", gliner_config_json) + + # When we're masquerading as `bert` for llama-quantize, mirror the required + # arch-scoped keys so its loader validates. llama.cpp's BERT loader expects + # u32 (not u64) for these fields; rgliner never reads them. + if post_quant is not None: + writer.add_metadata("bert.context_length", _U32(context_length)) + writer.add_metadata("bert.embedding_length", _U32(hidden_size)) + writer.add_metadata("bert.feed_forward_length", _U32(intermediate_size)) + writer.add_metadata("bert.block_count", _U32(num_layers)) + writer.add_metadata("bert.attention.head_count", _U32(num_heads)) + writer.add_metadata("bert.attention.layer_norm_epsilon", 1e-7) + + # Convert tensors + print(f"\nConverting {len(state_dict)} tensors to GGUF ({quantize.upper()})...") + + block_quant = quantize.startswith("q") + # Allow debugging by overriding via env var: + # GLINER_QUANT_INCLUDE="text.blk.0.ffn" -> only quantise tensors containing this substring + # GLINER_QUANT_EXCLUDE="token_embd" -> never quantise tensors containing this substring + quant_include = os.environ.get("GLINER_QUANT_INCLUDE") + quant_exclude = os.environ.get("GLINER_QUANT_EXCLUDE") + + for pytorch_name, tensor in state_dict.items(): + gguf_name = map_relex_weight_name(pytorch_name) + + try: + data = tensor.detach().float().cpu().numpy() + except RuntimeError: + import array + t = tensor.detach().float().cpu().contiguous() + data = np.frombuffer( + array.array('f', t.flatten().tolist()), + dtype=np.float32 + ).reshape(t.shape) + + logical_shape = tuple(data.shape) + tensor_type = default_ggml_type + + should_quant = block_quant + if should_quant and quant_include is not None and quant_include not in gguf_name: + should_quant = False + if should_quant and quant_exclude is not None and quant_exclude in gguf_name: + should_quant = False + + if should_quant: + # Row-wise block quantisation requires the inner dim to be a + # multiple of 32. 1-D tensors (biases, norms) and non-conforming + # tensors (odd vocab sizes etc.) stay at F32. + # + # Additional constraint: fusor's quantised matmul kernel produces + # NaN for all but the first output row when the output dimension + # is very small (the classifier head `scorer.out_mlp.3.weight` is + # [3, 3072]). Keep these tiny-output tensors at F32 - they're + # negligible bytes anyway. + MIN_OUT_DIM_FOR_QUANT = 32 + if ( + data.ndim >= 2 + and data.shape[-1] % 32 == 0 + and data.shape[0] >= MIN_OUT_DIM_FOR_QUANT + ): + data = _gguf_block_quant(data, quantize) + tensor_type = default_ggml_type + else: + data = np.ascontiguousarray(data, dtype=np.float32) + tensor_type = GGML_TYPE_F32 + else: + data = np.ascontiguousarray(data, dtype=np.float32) + tensor_type = GGML_TYPE_F32 + + print(f" {pytorch_name} -> {gguf_name} {logical_shape} [{_name_for_type(tensor_type)}]") + writer.add_tensor(gguf_name, data, tensor_type, shape=logical_shape) + + writer.write() + print(f"\nOutput: {output_path}") + print(f"Size: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB") + + if post_quant is not None: + print(f"\nQuantizing to {post_quant.upper()} via llama-quantize...") + try: + _llama_quantize(output_path, final_output_path, post_quant) + finally: + try: + os.remove(output_path) + except OSError: + pass + print(f"\nFinal output: {final_output_path}") + print(f"Final size: {os.path.getsize(final_output_path) / 1024 / 1024:.2f} MB") + + print("\nConversion complete!") + + +def main(): + parser = argparse.ArgumentParser(description="Convert GLiNER-RelEx PyTorch models to GGUF") + parser.add_argument( + "--model", "-m", + type=str, + required=True, + help="HuggingFace model ID (e.g., knowledgator/gliner-relex-multi-v1.0)" + ) + parser.add_argument( + "--output", "-o", + type=str, + required=True, + help="Output GGUF file path" + ) + parser.add_argument( + "--quantize", "-q", + type=str, + default="f32", + choices=sorted(QUANT_TYPES), + help=( + "Quantization type (default: f32). " + "All quantisation is done in-process via the `gguf` Python package. " + "k-quants (q4_k, q5_k, q6_k) aren't in this list - they are not " + "round-trip compatible with fusor's current quantised tensor reader." + ), + ) + parser.add_argument( + "--cache-dir", + type=str, + default=None, + help="HuggingFace cache directory" + ) + + args = parser.parse_args() + + convert_relex_to_gguf( + model_id=args.model, + output_path=args.output, + quantize=args.quantize, + cache_dir=args.cache_dir + ) + + +if __name__ == "__main__": + main() diff --git a/models/rgliner/scripts/convert_to_gguf.py b/models/rgliner/scripts/convert_to_gguf.py new file mode 100644 index 000000000..2ae1510d3 --- /dev/null +++ b/models/rgliner/scripts/convert_to_gguf.py @@ -0,0 +1,450 @@ +#!/usr/bin/env python3 +""" +Convert GLiNER PyTorch models to GGUF format. + +Usage: + python scripts/convert_to_gguf.py --model knowledgator/gliner-bi-edge-v2.0 --output gliner-bi-edge-v2.0.gguf + +This script converts both: +1. The main text encoder (ModernBERT/Ettin) + span layer weights +2. The label encoder projection weights (sentence transformer is loaded separately) +""" + +import argparse +import json +import os +import shutil +import struct +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch +from huggingface_hub import hf_hub_download, snapshot_download + +# GGUF/GGML constants, quant tables and the GGUFWriter live in the shared module. +from gguf_common import * # noqa: F401,F403 + + +def load_pytorch_model(model_id: str, cache_dir: str = None) -> Tuple[Dict[str, torch.Tensor], Dict]: + """Load PyTorch model from HuggingFace.""" + print(f"Downloading model: {model_id}") + + # Download the model files + model_dir = snapshot_download( + model_id, + cache_dir=cache_dir, + allow_patterns=["*.bin", "*.json", "*.safetensors"] + ) + + # Load config + config_path = os.path.join(model_dir, "gliner_config.json") + with open(config_path, 'r') as f: + config = json.load(f) + + # Load weights + weights_path = os.path.join(model_dir, "pytorch_model.bin") + if os.path.exists(weights_path): + print(f"Loading weights from: {weights_path}") + state_dict = torch.load(weights_path, map_location='cpu', weights_only=True) + else: + # Try safetensors + from safetensors.torch import load_file + weights_path = os.path.join(model_dir, "model.safetensors") + print(f"Loading weights from: {weights_path}") + state_dict = load_file(weights_path) + + return state_dict, config + + +def map_weight_name(pytorch_name: str) -> str: + """Map PyTorch weight names to GGUF conventions.""" + name = pytorch_name + + # ===== Encoder output projection ===== + # Some bi-encoder variants (small/base/large v2.0) project the text-encoder + # hidden state down to the shared label-aligned dim (e.g. 512 -> 384). + name = name.replace("token_rep_layer.projection.weight", "text.output_proj.weight") + name = name.replace("token_rep_layer.projection.bias", "text.output_proj.bias") + + # ===== Text Encoder (ModernBERT/Ettin) ===== + # token_rep_layer.bert_layer.model.embeddings.tok_embeddings.weight -> text.token_embd.weight + name = name.replace("token_rep_layer.bert_layer.model.embeddings.tok_embeddings.weight", "text.token_embd.weight") + name = name.replace("token_rep_layer.bert_layer.model.embeddings.norm.weight", "text.embd_norm.weight") + + # token_rep_layer.bert_layer.model.layers.X -> text.blk.X + name = name.replace("token_rep_layer.bert_layer.model.layers.", "text.blk.") + name = name.replace("token_rep_layer.bert_layer.model.final_norm.weight", "text.output_norm.weight") + + # ModernBERT attention (fused Wqkv) + name = name.replace(".attn.Wqkv.", ".attn_qkv.") + name = name.replace(".attn.Wo.", ".attn_output.") + + # ModernBERT FFN (GeGLU with fused Wi) + name = name.replace(".mlp.Wi.", ".ffn_gate_up.") # Fused gate+up + name = name.replace(".mlp.Wo.", ".ffn_down.") + name = name.replace(".mlp_norm.", ".ffn_norm.") + + # ===== Label Encoder (BERT/MiniLM) ===== + name = name.replace("token_rep_layer.labels_encoder.model.", "label.") + + # BERT embeddings (rbert expects token_types, token_embd_norm) + name = name.replace("label.embeddings.word_embeddings.", "label.token_embd.") + name = name.replace("label.embeddings.position_embeddings.", "label.position_embd.") + name = name.replace("label.embeddings.token_type_embeddings.", "label.token_types.") + name = name.replace("label.embeddings.LayerNorm.", "label.token_embd_norm.") + + # BERT layers + name = name.replace("label.encoder.layer.", "label.blk.") + + # BERT attention (rbert uses attn_output_norm, not attn_norm) + name = name.replace(".attention.self.query.", ".attn_q.") + name = name.replace(".attention.self.key.", ".attn_k.") + name = name.replace(".attention.self.value.", ".attn_v.") + name = name.replace(".attention.output.dense.", ".attn_output.") + name = name.replace(".attention.output.LayerNorm.", ".attn_output_norm.") + + # BERT FFN (rbert uses layer_output_norm, not ffn_norm) + name = name.replace(".intermediate.dense.", ".ffn_up.") + name = name.replace(".output.dense.", ".ffn_down.") + name = name.replace(".output.LayerNorm.", ".layer_output_norm.") + + # BERT pooler + name = name.replace("label.pooler.dense.", "label.pooler.") + + # ===== BiLSTM ===== + # Keep rnn.lstm.* as-is for now + name = name.replace("rnn.lstm.", "rnn.") + + # ===== Span Representation Layer ===== + name = name.replace("span_rep_layer.span_rep_layer.project_start.0.", "span.start_fc1.") + name = name.replace("span_rep_layer.span_rep_layer.project_start.3.", "span.start_fc2.") + name = name.replace("span_rep_layer.span_rep_layer.project_end.0.", "span.end_fc1.") + name = name.replace("span_rep_layer.span_rep_layer.project_end.3.", "span.end_fc2.") + name = name.replace("span_rep_layer.span_rep_layer.out_project.0.", "span.out_fc1.") + name = name.replace("span_rep_layer.span_rep_layer.out_project.3.", "span.out_fc2.") + + # ===== Prompt/Label Projection ===== + name = name.replace("prompt_rep_layer.0.", "label_proj.0.") + name = name.replace("prompt_rep_layer.3.", "label_proj.2.") + + return name + + +def _llama_quantize(f32_path: str, output_path: str, quant_type: str, keep_f32: List[str]) -> None: + """Run `llama-quantize` to convert an f32 GGUF to a k-quant type. + + `keep_f32` is a list of tensor names that should remain at F32 (tiny + classifier heads hit a fusor bug; see the gliner-relex project notes). + """ + binary = shutil.which("llama-quantize") + if binary is None: + raise RuntimeError( + "`llama-quantize` not found in PATH. Install llama.cpp to use k-quant types:\n" + " brew install llama.cpp # macOS\n" + " or build from https://github.com/ggml-org/llama.cpp" + ) + cmd = [binary] + for name in keep_f32: + cmd += ["--tensor-type", f"{name}=f32"] + cmd += [f32_path, output_path, quant_type.upper()] + print(f"\n$ {' '.join(cmd)}") + result = subprocess.run(cmd) + if result.returncode != 0: + raise RuntimeError( + f"llama-quantize failed (exit {result.returncode}) for {quant_type!r}. " + f"See output above. The f32 GGUF was kept at {f32_path} for inspection." + ) + + +def _quantize_tensor(data: np.ndarray, quant: str, default_ggml_type: int) -> Tuple[np.ndarray, int, Tuple[int, ...]]: + """Quantise a single tensor, falling back to F32 for unsupported shapes. + + Returns (packed_data, ggml_type, logical_shape). Applies the same rules as + the relex converter: only 2-D+ tensors with inner dim divisible by 32 and + outer dim >= 32 are block-quantised. The lower outer-dim threshold avoids + a fusor quantised-matmul bug on tiny classifier heads. + """ + logical_shape = tuple(data.shape) + MIN_OUT_DIM_FOR_QUANT = 32 + if ( + quant.startswith("q") + and data.ndim >= 2 + and data.shape[-1] % 32 == 0 + and data.shape[0] >= MIN_OUT_DIM_FOR_QUANT + ): + return _gguf_block_quant(data, quant), default_ggml_type, logical_shape + return np.ascontiguousarray(data, dtype=np.float32), GGML_TYPE_F32, logical_shape + + +def convert_gliner_to_gguf( + model_id: str, + output_path: str, + quantize: str = "f32", + cache_dir: str = None +): + """Convert GLiNER model to GGUF format. + + Creates two GGUF files: + - {output_path}: Main model (text encoder, span layer, projection) + - {output_path_stem}-label-encoder.gguf: Label encoder (BERT/MiniLM) + """ + + # Load model + state_dict, config = load_pytorch_model(model_id, cache_dir) + + # Print model structure + print("\nModel weights:") + for name, tensor in state_dict.items(): + print(f" {name}: {tensor.shape} {tensor.dtype}") + + # Determine quantization type and post-processing + quantize = quantize.lower() + if quantize not in QUANT_TYPES: + raise ValueError( + f"Unsupported quantization: {quantize!r}. Supported: {', '.join(sorted(QUANT_TYPES))}" + ) + post_quant = None + if quantize in LLAMA_QUANT_TYPES: + post_quant = quantize + quantize = "f32" + default_ggml_type = _ggml_type_for(quantize) + + # When post-quantising, write the raw f32 file to a temp path first. + main_label_output = output_path.replace(".gguf", "-label-encoder.gguf") + final_main = output_path + final_label = main_label_output + if post_quant is not None: + main_tmp = tempfile.NamedTemporaryFile( + prefix=os.path.basename(output_path).rsplit(".", 1)[0] + ".f32.", + suffix=".gguf", + delete=False, + dir=os.path.dirname(os.path.abspath(output_path)) or None, + ) + label_tmp = tempfile.NamedTemporaryFile( + prefix=os.path.basename(main_label_output).rsplit(".", 1)[0] + ".f32.", + suffix=".gguf", + delete=False, + dir=os.path.dirname(os.path.abspath(main_label_output)) or None, + ) + output_path = main_tmp.name + main_label_output = label_tmp.name + main_tmp.close() + label_tmp.close() + + # Separate label encoder weights from main model weights + label_encoder_weights = {} + main_model_weights = {} + + for pytorch_name, tensor in state_dict.items(): + if "token_rep_layer.labels_encoder" in pytorch_name: + label_encoder_weights[pytorch_name] = tensor + else: + main_model_weights[pytorch_name] = tensor + + # ============ Main Model GGUF ============ + writer = GGUFWriter(output_path) + + # Add metadata. Masquerade as `bert` when we'll feed this to llama-quantize + # so its loader accepts the custom architecture. rgliner reads `gliner.*` + # keys, not `general.architecture`. + main_arch = "bert" if post_quant is not None else "gliner" + writer.add_metadata("general.architecture", main_arch) + writer.add_metadata("general.name", model_id.split("/")[-1]) + writer.add_metadata("general.quantization_version", 2) + + # GLiNER-specific metadata + writer.add_metadata("gliner.max_width", config.get("max_width", 12)) + writer.add_metadata("gliner.span_mode", config.get("span_mode", "markerV0")) + writer.add_metadata("gliner.subtoken_pooling", config.get("subtoken_pooling", "first")) + + # Encoder config - use standard GGUF naming for model loading + encoder_config = config.get("encoder_config", {}) + hidden_size = encoder_config.get("hidden_size", 384) + num_heads = encoder_config.get("num_attention_heads", 6) + num_layers = encoder_config.get("num_hidden_layers", 10) + intermediate_size = encoder_config.get("intermediate_size", 576) + vocab_size = encoder_config.get("vocab_size", 50368) + context_length = encoder_config.get("max_position_embeddings", 8192) + # ModernBERT/Ettin alternate full-global attention every Nth layer with + # sliding-window local attention on the rest, and the two layer types may use + # different RoPE bases. Defaults match ModernBERT-base. + global_rope_theta = encoder_config.get("global_rope_theta", 160000.0) + local_rope_theta = encoder_config.get("local_rope_theta", global_rope_theta) + global_attn_every_n_layers = encoder_config.get("global_attn_every_n_layers", 3) + local_attention = encoder_config.get("local_attention", 128) + + # Standard GGUF metadata (without architecture prefix - the loader adds it) + writer.add_metadata("gliner.attention.head_count", num_heads) + writer.add_metadata("gliner.attention.head_count_kv", num_heads) # No GQA in this model + writer.add_metadata("gliner.block_count", num_layers) + writer.add_metadata("gliner.embedding_length", hidden_size) + writer.add_metadata("gliner.feed_forward_length", intermediate_size) + writer.add_metadata("gliner.context_length", context_length) + writer.add_metadata("gliner.rope.freq_base", float(global_rope_theta)) + writer.add_metadata("gliner.rope.local_freq_base", float(local_rope_theta)) + writer.add_metadata( + "gliner.attention.global_attn_every_n_layers", global_attn_every_n_layers + ) + writer.add_metadata("gliner.attention.local_attention", local_attention) + writer.add_metadata("gliner.attention.layer_norm_rms_epsilon", 1e-5) + writer.add_metadata("gliner.vocab_size", vocab_size) + + # When feeding to llama-quantize, mirror arch keys with u32 scalars. + if post_quant is not None: + writer.add_metadata("bert.context_length", _U32(context_length)) + writer.add_metadata("bert.embedding_length", _U32(hidden_size)) + writer.add_metadata("bert.feed_forward_length", _U32(intermediate_size)) + writer.add_metadata("bert.block_count", _U32(num_layers)) + writer.add_metadata("bert.attention.head_count", _U32(num_heads)) + writer.add_metadata("bert.attention.layer_norm_epsilon", 1e-5) + + # Convert main model tensors + print(f"\nConverting {len(main_model_weights)} main model tensors to GGUF ({quantize.upper()})...") + + for pytorch_name, tensor in main_model_weights.items(): + gguf_name = map_weight_name(pytorch_name) + + # Convert to numpy (workaround for numpy/torch incompatibility) + try: + data = tensor.detach().float().cpu().numpy() + except RuntimeError: + import array + t = tensor.detach().float().cpu().contiguous() + data = np.frombuffer( + array.array('f', t.flatten().tolist()), + dtype=np.float32 + ).reshape(t.shape) + + packed, tensor_type, logical_shape = _quantize_tensor(data, quantize, default_ggml_type) + print(f" {pytorch_name} -> {gguf_name} {logical_shape} [{_name_for_type(tensor_type)}]") + writer.add_tensor(gguf_name, packed, tensor_type, shape=logical_shape) + + writer.write() + print(f"Main model output: {output_path}") + print(f"Size: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB") + + # ============ Label Encoder GGUF ============ + # Create separate file for label encoder (without prefix, for rbert compatibility) + label_output_path = main_label_output + label_writer = GGUFWriter(label_output_path) + + # Add BERT metadata + labels_config = config.get("labels_encoder_config", {}) + label_writer.add_metadata("general.architecture", "bert") + label_writer.add_metadata("general.name", model_id.split("/")[-1] + "-label-encoder") + + label_hidden = labels_config.get("hidden_size", 384) + label_heads = labels_config.get("num_attention_heads", 12) + label_layers = labels_config.get("num_hidden_layers", 6) + label_intermediate = labels_config.get("intermediate_size", 1536) + label_vocab = labels_config.get("vocab_size", 30522) + label_max_pos = labels_config.get("max_position_embeddings", 512) + + # Use u32 for arch-scoped ints when we'll feed this through llama-quantize. + int_ctor = _U32 if post_quant is not None else (lambda x: x) + label_writer.add_metadata("bert.attention.head_count", int_ctor(label_heads)) + label_writer.add_metadata("bert.block_count", int_ctor(label_layers)) + label_writer.add_metadata("bert.embedding_length", int_ctor(label_hidden)) + label_writer.add_metadata("bert.feed_forward_length", int_ctor(label_intermediate)) + label_writer.add_metadata("bert.context_length", int_ctor(label_max_pos)) + label_writer.add_metadata("bert.attention.layer_norm_epsilon", 1e-12) + label_writer.add_metadata("bert.vocab_size", int_ctor(label_vocab)) + + # Convert label encoder tensors (remove prefix so rbert can load them) + print(f"\nConverting {len(label_encoder_weights)} label encoder tensors to GGUF ({quantize.upper()})...") + + for pytorch_name, tensor in label_encoder_weights.items(): + # Map name but remove the "label." prefix for rbert compatibility + gguf_name = map_weight_name(pytorch_name) + if gguf_name.startswith("label."): + gguf_name = gguf_name[6:] # Remove "label." prefix + + try: + data = tensor.detach().float().cpu().numpy() + except RuntimeError: + import array + t = tensor.detach().float().cpu().contiguous() + data = np.frombuffer( + array.array('f', t.flatten().tolist()), + dtype=np.float32 + ).reshape(t.shape) + + packed, tensor_type, logical_shape = _quantize_tensor(data, quantize, default_ggml_type) + print(f" {pytorch_name} -> {gguf_name} {logical_shape} [{_name_for_type(tensor_type)}]") + label_writer.add_tensor(gguf_name, packed, tensor_type, shape=logical_shape) + + label_writer.write() + print(f"Label encoder output: {label_output_path}") + print(f"Size: {os.path.getsize(label_output_path) / 1024 / 1024:.2f} MB") + + if post_quant is not None: + print(f"\nQuantizing main model to {post_quant.upper()} via llama-quantize...") + _llama_quantize(output_path, final_main, post_quant, keep_f32=[]) + try: + os.remove(output_path) + except OSError: + pass + print(f"Final main output: {final_main}") + print(f"Final main size: {os.path.getsize(final_main) / 1024 / 1024:.2f} MB") + + print(f"\nQuantizing label encoder to {post_quant.upper()} via llama-quantize...") + _llama_quantize(label_output_path, final_label, post_quant, keep_f32=[]) + try: + os.remove(label_output_path) + except OSError: + pass + print(f"Final label output: {final_label}") + print(f"Final label size: {os.path.getsize(final_label) / 1024 / 1024:.2f} MB") + + print(f"\nConversion complete!") + + +def main(): + parser = argparse.ArgumentParser(description="Convert GLiNER PyTorch models to GGUF") + parser.add_argument( + "--model", "-m", + type=str, + required=True, + help="HuggingFace model ID (e.g., knowledgator/gliner-bi-edge-v2.0)" + ) + parser.add_argument( + "--output", "-o", + type=str, + required=True, + help="Output GGUF file path" + ) + parser.add_argument( + "--quantize", "-q", + type=str, + default="f32", + choices=sorted(QUANT_TYPES), + help=( + "Quantization type (default: f32). " + "Block quants (q4_0, q5_0, q8_0, ...) are packed in-process via the " + "`gguf` Python package. K-quants (q4_k, q5_k, q6_k, ...) require " + "`llama-quantize` in PATH and are applied as a post-processing step." + ), + ) + parser.add_argument( + "--cache-dir", + type=str, + default=None, + help="HuggingFace cache directory" + ) + + args = parser.parse_args() + + convert_gliner_to_gguf( + model_id=args.model, + output_path=args.output, + quantize=args.quantize, + cache_dir=args.cache_dir + ) + + +if __name__ == "__main__": + main() diff --git a/models/rgliner/scripts/gguf_common.py b/models/rgliner/scripts/gguf_common.py new file mode 100644 index 000000000..7bb4cba45 --- /dev/null +++ b/models/rgliner/scripts/gguf_common.py @@ -0,0 +1,280 @@ +"""Shared GGUF-writing infrastructure for the GLiNER converters. + +Both `convert_to_gguf.py` (bi-encoder) and `convert_relex_to_gguf.py` (RelEx) +emit GGUF files with the byte layout fusor expects. This module holds the +quantisation tables, GGUF/GGML constants, and the `GGUFWriter` they share; each +converter keeps its own model-specific weight mapping and `main()`. +""" + +import struct +from typing import Any, Dict, List, Tuple + +import numpy as np + +__all__ = [ + "IN_PROCESS_QUANTS", + "LLAMA_QUANT_TYPES", + "QUANT_TYPES", + "GGUF_MAGIC", + "GGUF_VERSION", + "GGUF_TYPE_UINT8", + "GGUF_TYPE_INT8", + "GGUF_TYPE_UINT16", + "GGUF_TYPE_INT16", + "GGUF_TYPE_UINT32", + "GGUF_TYPE_INT32", + "GGUF_TYPE_FLOAT32", + "GGUF_TYPE_BOOL", + "GGUF_TYPE_STRING", + "GGUF_TYPE_ARRAY", + "GGUF_TYPE_UINT64", + "GGUF_TYPE_INT64", + "GGUF_TYPE_FLOAT64", + "GGML_TYPE_F32", + "GGML_TYPE_F16", + "GGML_TYPE_Q4_0", + "GGML_TYPE_Q4_1", + "GGML_TYPE_Q5_0", + "GGML_TYPE_Q5_1", + "GGML_TYPE_Q8_0", + "GGML_TYPE_Q8_1", + "GGML_TYPE_BF16", + "_ggml_type_for", + "_name_for_type", + "_gguf_block_quant", + "_U32", + "GGUFWriter", +] + + +# Quantization targets. In-process ones are packed directly by `gguf.quants.quantize`; +# k-quants are produced by shelling out to `llama-quantize` post-hoc. +IN_PROCESS_QUANTS = { + "f32", "f16", "bf16", + "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", +} +LLAMA_QUANT_TYPES = { + "q2_k", "q3_k", "q3_k_s", "q3_k_m", "q3_k_l", + "q4_k", "q4_k_s", "q4_k_m", + "q5_k", "q5_k_s", "q5_k_m", + "q6_k", +} +QUANT_TYPES = IN_PROCESS_QUANTS | LLAMA_QUANT_TYPES + + +# GGUF constants +GGUF_MAGIC = 0x46554747 # "GGUF" in little-endian +GGUF_VERSION = 3 + +# GGUF data types +GGUF_TYPE_UINT8 = 0 +GGUF_TYPE_INT8 = 1 +GGUF_TYPE_UINT16 = 2 +GGUF_TYPE_INT16 = 3 +GGUF_TYPE_UINT32 = 4 +GGUF_TYPE_INT32 = 5 +GGUF_TYPE_FLOAT32 = 6 +GGUF_TYPE_BOOL = 7 +GGUF_TYPE_STRING = 8 +GGUF_TYPE_ARRAY = 9 +GGUF_TYPE_UINT64 = 10 +GGUF_TYPE_INT64 = 11 +GGUF_TYPE_FLOAT64 = 12 + +# GGML tensor types +GGML_TYPE_F32 = 0 +GGML_TYPE_F16 = 1 +GGML_TYPE_Q4_0 = 2 +GGML_TYPE_Q4_1 = 3 +GGML_TYPE_Q5_0 = 6 +GGML_TYPE_Q5_1 = 7 +GGML_TYPE_Q8_0 = 8 +GGML_TYPE_Q8_1 = 9 +GGML_TYPE_BF16 = 30 + + +def _ggml_type_for(quant: str) -> int: + return { + "f32": GGML_TYPE_F32, + "f16": GGML_TYPE_F16, + "bf16": GGML_TYPE_BF16, + "q4_0": GGML_TYPE_Q4_0, + "q4_1": GGML_TYPE_Q4_1, + "q5_0": GGML_TYPE_Q5_0, + "q5_1": GGML_TYPE_Q5_1, + "q8_0": GGML_TYPE_Q8_0, + }[quant] + + +def _name_for_type(ggml_type: int) -> str: + return { + GGML_TYPE_F32: "F32", + GGML_TYPE_F16: "F16", + GGML_TYPE_BF16: "BF16", + GGML_TYPE_Q4_0: "Q4_0", + GGML_TYPE_Q4_1: "Q4_1", + GGML_TYPE_Q5_0: "Q5_0", + GGML_TYPE_Q5_1: "Q5_1", + GGML_TYPE_Q8_0: "Q8_0", + }.get(ggml_type, f"type={ggml_type}") + + +def _gguf_block_quant(data: np.ndarray, quant: str) -> np.ndarray: + """Return the packed byte array for a block-quantised tensor. + + Uses `gguf.quants.quantize` (Python package from llama.cpp). The inner row + dimension must be a multiple of 32 for q4_0/q5_0/q8_0, else we fall back to + f16 for that tensor (caller decides). + """ + import gguf # Deferred import - only needed for block quants. + qtype = { + "q4_0": gguf.GGMLQuantizationType.Q4_0, + "q4_1": gguf.GGMLQuantizationType.Q4_1, + "q5_0": gguf.GGMLQuantizationType.Q5_0, + "q5_1": gguf.GGMLQuantizationType.Q5_1, + "q8_0": gguf.GGMLQuantizationType.Q8_0, + }[quant] + return gguf.quants.quantize(data, qtype) + + +class _U32(int): + """Force a metadata int to be written as GGUF u32 (required by llama-quantize's arch loader).""" + + +class GGUFWriter: + """Simple GGUF file writer.""" + + def __init__(self, path: str): + self.path = path + self.metadata: Dict[str, Any] = {} + # (name, data, ggml_type, logical_shape) + self.tensors: List[Tuple[str, np.ndarray, int, Tuple[int, ...]]] = [] + + def add_metadata(self, key: str, value: Any): + """Add metadata key-value pair.""" + self.metadata[key] = value + + def add_tensor( + self, + name: str, + data: np.ndarray, + ggml_type: int = GGML_TYPE_F32, + shape: Tuple[int, ...] = None, + ): + """Add a tensor. `shape` is the logical (un-packed) shape; defaults to data.shape.""" + if shape is None: + shape = tuple(data.shape) + self.tensors.append((name, data, ggml_type, tuple(shape))) + + def _write_string(self, f, s: str): + """Write a GGUF string (length-prefixed UTF-8).""" + encoded = s.encode('utf-8') + f.write(struct.pack('> 16) & 0xFFFF).astype(np.uint16) + + # Write tensor info + # GGUF stores dimensions in reverse order (column-major) + # Reader reverses them back, so we write reversed to get original order + self._write_string(f, name) + f.write(struct.pack(', + + /// Maximum span width in words (default: 12) + #[serde(default = "default_max_width")] + pub max_width: usize, +} + +fn default_max_width() -> usize { + 12 +} + +impl GlinerConfig { + /// Parse config from JSON bytes. + pub fn from_json(json: &[u8]) -> Result { + serde_json::from_slice(json) + } + + /// Whether the tokenizer should add [CLS]/[SEP] special tokens around text. + /// + /// Matches Python GLiNER's `_set_tokenizer_spec_tokens` behavior: + /// ModernBERT/ettin-style encoders (which have `add_bos_token=False` + /// semantics because they have no bos token) are fed raw text without + /// [CLS]/[SEP] wrappers. DeBERTa/RoBERTa/XLM-R family keep them. + pub fn should_add_special_tokens(&self) -> bool { + match self.model_name.as_deref() { + Some(name) => { + let lower = name.to_ascii_lowercase(); + !(lower.contains("ettin") + || lower.contains("modernbert") + || lower.contains("modern-bert")) + } + None => true, + } + } +} diff --git a/models/rgliner/src/decoding.rs b/models/rgliner/src/decoding.rs new file mode 100644 index 000000000..b1220bae6 --- /dev/null +++ b/models/rgliner/src/decoding.rs @@ -0,0 +1,370 @@ +//! Entity decoding with flat and nested NER support. + +/// A recognized named entity. +#[derive(Debug, Clone)] +pub struct Entity { + /// The entity text span. + pub text: String, + /// The entity label/type. + pub label: String, + /// Start character offset in the original text. + pub start_char: usize, + /// End character offset in the original text (exclusive). + pub end_char: usize, + /// Start word index. + pub start_word: usize, + /// End word index (inclusive). + pub end_word: usize, + /// Confidence score (0.0 to 1.0). + pub score: f32, +} + +/// Decoding mode for NER. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum DecodingMode { + /// Flat NER: no overlapping entities allowed. + /// Uses greedy non-maximum suppression. + #[default] + Flat, + /// Nested NER: overlapping entities allowed if one fully contains the other. + /// Partial overlaps are still forbidden. + Nested, +} + +/// Entity decoder using non-maximum suppression. +pub struct Decoder { + /// Confidence threshold for entity detection. + threshold: f32, + /// Decoding mode (flat or nested). + mode: DecodingMode, +} + +impl Decoder { + /// Create a new decoder with the given threshold and mode. + pub fn new(threshold: f32, mode: DecodingMode) -> Self { + Self { threshold, mode } + } + + /// Set the confidence threshold. + pub fn with_threshold(mut self, threshold: f32) -> Self { + self.threshold = threshold; + self + } + + /// Set the decoding mode. + pub fn with_mode(mut self, mode: DecodingMode) -> Self { + self.mode = mode; + self + } + + /// Decode entity predictions for a single text. + /// + /// # Arguments + /// * `scores` - Score matrix [num_spans, num_labels] after sigmoid + /// * `span_indices` - (start_word, end_word) for each span + /// * `word_offsets` - (start_char, end_char) for each word + /// * `labels` - Label strings + /// * `text` - Original input text + pub fn decode( + &self, + scores: &[f32], + num_labels: usize, + span_indices: &[(usize, usize)], + word_offsets: &[(usize, usize)], + labels: &[&str], + text: &str, + ) -> Vec { + // Collect all predictions above threshold + let mut candidates: Vec<(usize, usize, f32)> = Vec::new(); // (span_idx, label_idx, score) + + let num_spans = span_indices.len(); + for span_idx in 0..num_spans { + for label_idx in 0..num_labels { + let score = scores[span_idx * num_labels + label_idx]; + if score >= self.threshold { + candidates.push((span_idx, label_idx, score)); + } + } + } + + // Sort by score descending + candidates.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal)); + + match self.mode { + DecodingMode::Flat => { + self.decode_flat(candidates, span_indices, word_offsets, labels, text) + } + DecodingMode::Nested => { + self.decode_nested(candidates, span_indices, word_offsets, labels, text) + } + } + } + + fn decode_flat( + &self, + candidates: Vec<(usize, usize, f32)>, + span_indices: &[(usize, usize)], + word_offsets: &[(usize, usize)], + labels: &[&str], + text: &str, + ) -> Vec { + let num_words = word_offsets.len(); + let mut entities = Vec::new(); + let mut used_positions: Vec = vec![false; num_words]; + + for (span_idx, label_idx, score) in candidates { + let (start_word, end_word) = span_indices[span_idx]; + + // Check if any word in span is already used + let overlaps = + (start_word..=end_word).any(|w| used_positions.get(w).copied().unwrap_or(false)); + + if !overlaps { + // Mark words as used + for w in start_word..=end_word { + if w < used_positions.len() { + used_positions[w] = true; + } + } + + if let Some(entity) = Self::create_entity( + start_word, + end_word, + label_idx, + score, + word_offsets, + labels, + text, + ) { + entities.push(entity); + } + } + } + + // Sort by position + entities.sort_by_key(|e| e.start_char); + entities + } + + fn decode_nested( + &self, + candidates: Vec<(usize, usize, f32)>, + span_indices: &[(usize, usize)], + word_offsets: &[(usize, usize)], + labels: &[&str], + text: &str, + ) -> Vec { + let mut entities = Vec::new(); + // Track selected (start, end, label) triples to avoid exact duplicates + let mut selected: Vec<(usize, usize, usize)> = Vec::new(); + + for (span_idx, label_idx, score) in candidates { + let (start_word, end_word) = span_indices[span_idx]; + let key = (start_word, end_word, label_idx); + + // Check for partial overlap with already selected entities + let has_partial_overlap = selected.iter().any(|(sel_start, sel_end, _)| { + self.is_partial_overlap(start_word, end_word, *sel_start, *sel_end) + }); + + // For nested NER, allow if no partial overlap and not exact duplicate + if !has_partial_overlap && !selected.contains(&key) { + selected.push(key); + + if let Some(entity) = Self::create_entity( + start_word, + end_word, + label_idx, + score, + word_offsets, + labels, + text, + ) { + entities.push(entity); + } + } + } + + // Sort by position, then by span length descending (outer spans first) + entities.sort_by(|a, b| { + a.start_char + .cmp(&b.start_char) + .then_with(|| b.end_char.cmp(&a.end_char)) + }); + entities + } + + /// Check if two spans have partial overlap (overlap but neither contains the other). + fn is_partial_overlap(&self, start1: usize, end1: usize, start2: usize, end2: usize) -> bool { + // Check if spans overlap at all + let overlaps = start1 <= end2 && start2 <= end1; + if !overlaps { + return false; + } + + // Check if one fully contains the other (not partial overlap) + let one_contains_other = + (start1 <= start2 && end1 >= end2) || (start2 <= start1 && end2 >= end1); + + overlaps && !one_contains_other + } + + fn create_entity( + start_word: usize, + end_word: usize, + label_idx: usize, + score: f32, + word_offsets: &[(usize, usize)], + labels: &[&str], + text: &str, + ) -> Option { + if start_word >= word_offsets.len() || end_word >= word_offsets.len() { + return None; + } + if label_idx >= labels.len() { + return None; + } + + let start_char = word_offsets[start_word].0; + let end_char = word_offsets[end_word].1; + + if end_char > text.len() { + return None; + } + + Some(Entity { + text: text[start_char..end_char].to_string(), + label: labels[label_idx].to_string(), + start_char, + end_char, + start_word, + end_word, + score, + }) + } +} + +impl Default for Decoder { + fn default() -> Self { + Self { + threshold: 0.5, + mode: DecodingMode::default(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const TEXT: &str = "Barack Obama visited Hawaii"; + // word -> (start_char, end_char): Barack, Obama, visited, Hawaii + fn word_offsets() -> Vec<(usize, usize)> { + vec![(0, 6), (7, 12), (13, 20), (21, 27)] + } + + /// Flatten `(span_idx, label_idx, score)` triples into the row-major + /// `[num_spans * num_labels]` score buffer `decode` expects. + fn scores(num_spans: usize, num_labels: usize, entries: &[(usize, usize, f32)]) -> Vec { + let mut s = vec![0.0; num_spans * num_labels]; + for &(span, label, score) in entries { + s[span * num_labels + label] = score; + } + s + } + + #[test] + fn flat_suppresses_overlapping_lower_scoring_span() { + // span 0 = "Barack Obama" (words 0..=1), span 1 = "Barack" (0..=0), + // span 2 = "Hawaii" (3..=3). Labels: 0=person, 1=location. + let span_indices = [(0, 1), (0, 0), (3, 3)]; + let labels = ["person", "location"]; + let s = scores( + 3, + 2, + &[ + (0, 0, 0.9), // "Barack Obama" / person + (1, 0, 0.6), // "Barack" / person (overlaps span 0) + (2, 1, 0.8), // "Hawaii" / location + ], + ); + + let entities = Decoder::new(0.5, DecodingMode::Flat).decode( + &s, + 2, + &span_indices, + &word_offsets(), + &labels, + TEXT, + ); + + // span 1 shares word 0 with the higher-scoring span 0 -> suppressed. + assert_eq!(entities.len(), 2); + assert_eq!(entities[0].text, "Barack Obama"); + assert_eq!(entities[0].label, "person"); + assert_eq!((entities[0].start_char, entities[0].end_char), (0, 12)); + assert_eq!(entities[1].text, "Hawaii"); + assert_eq!(entities[1].label, "location"); + assert_eq!((entities[1].start_char, entities[1].end_char), (21, 27)); + } + + #[test] + fn nested_keeps_contained_span_but_drops_partial_overlap() { + // span 0 = (0..=1), span 1 = (0..=0) fully contained in span 0, + // span 2 = (1..=2) partially overlaps span 0. + let span_indices = [(0, 1), (0, 0), (1, 2)]; + let labels = ["person"]; + let s = scores(3, 1, &[(0, 0, 0.9), (1, 0, 0.7), (2, 0, 0.6)]); + + let entities = Decoder::new(0.5, DecodingMode::Nested).decode( + &s, + 1, + &span_indices, + &word_offsets(), + &labels, + TEXT, + ); + + // Contained span (0..=0) is allowed; partial overlap (1..=2) is not. + assert_eq!(entities.len(), 2); + let spans: Vec<(usize, usize)> = entities + .iter() + .map(|e| (e.start_word, e.end_word)) + .collect(); + assert!(spans.contains(&(0, 1))); + assert!(spans.contains(&(0, 0))); + assert!(!spans.contains(&(1, 2))); + } + + #[test] + fn scores_below_threshold_are_dropped() { + let span_indices = [(0, 0)]; + let labels = ["person"]; + let s = scores(1, 1, &[(0, 0, 0.49)]); + let entities = Decoder::new(0.5, DecodingMode::Flat).decode( + &s, + 1, + &span_indices, + &word_offsets(), + &labels, + TEXT, + ); + assert!(entities.is_empty()); + } + + #[test] + fn is_partial_overlap_predicate() { + let d = Decoder::default(); + // Disjoint. + assert!(!d.is_partial_overlap(0, 1, 2, 3)); + // Adjacent ranges sharing word 1, neither containing the other -> partial. + assert!(d.is_partial_overlap(0, 1, 1, 2)); + // Containment is not "partial". + assert!(!d.is_partial_overlap(0, 3, 1, 2)); + assert!(!d.is_partial_overlap(1, 2, 0, 3)); + // Genuine partial overlap. + assert!(d.is_partial_overlap(0, 2, 1, 3)); + // Identical spans -> containment, not partial. + assert!(!d.is_partial_overlap(1, 2, 1, 2)); + } +} diff --git a/models/rgliner/src/error.rs b/models/rgliner/src/error.rs new file mode 100644 index 000000000..bf2f5382b --- /dev/null +++ b/models/rgliner/src/error.rs @@ -0,0 +1,44 @@ +//! Error types for rgliner. + +use kalosm_common::CacheError; + +/// An error that can occur when loading a GLiNER model. +#[derive(Debug, thiserror::Error)] +pub enum GlinerLoadingError { + /// An error that can occur when trying to download model files. + #[error("Failed to download model files: {0}")] + DownloadingError(#[from] CacheError), + /// An error that can occur when trying to load the model. + #[error("Failed to load model: {0}")] + LoadModel(#[from] fusor::Error), + /// An error that can occur when trying to load the tokenizer. + #[error("Failed to load tokenizer: {0}")] + LoadTokenizer(tokenizers::Error), + /// An error that can occur when trying to load the config. + #[error("Failed to load config: {0}")] + LoadConfig(serde_json::Error), + /// A required GLiNER-RelEx special token is absent from the tokenizer. + #[error( + "Tokenizer is missing required GLiNER-RelEx special token(s): {0}. The model's embedded \ + tokenizer is for the wrong variant or is corrupt — re-run the conversion script or supply \ + a matching tokenizer via `GlinerRelExSource::with_tokenizer`." + )] + MissingSpecialToken(String), + /// Label encoder loading error. + #[error("Failed to load label encoder: {0}")] + LabelEncoder(#[from] rbert::BertLoadingError), +} + +/// An error that can occur when running GLiNER inference. +#[derive(Debug, thiserror::Error)] +pub enum GlinerError { + /// An error that can occur when running tensor operations. + #[error("Tensor operation error: {0}")] + Fusor(#[from] fusor::Error), + /// An error that can occur when tokenizing text. + #[error("Tokenization error: {0}")] + Tokenizer(tokenizers::Error), + /// An error that can occur with the label encoder. + #[error("Label encoder error: {0}")] + LabelEncoder(#[from] rbert::BertError), +} diff --git a/models/rgliner/src/lib.rs b/models/rgliner/src/lib.rs new file mode 100644 index 000000000..ea5b9e7a5 --- /dev/null +++ b/models/rgliner/src/lib.rs @@ -0,0 +1,559 @@ +//! # rgliner +//! +//! GLiNER bi-encoder Named Entity Recognition for Rust. +//! +//! GLiNER (Generalist Lightweight Model for Named Entity Recognition) identifies +//! arbitrary entity types at inference time using natural language labels. +//! +//! ## Usage +//! +//! ```rust, no_run +//! use rgliner::*; +//! +//! #[tokio::main] +//! async fn main() -> anyhow::Result<()> { +//! let mut gliner = Gliner::new().await?; +//! +//! let labels = ["person", "organization", "location"]; +//! let text = "Apple Inc. was founded by Steve Jobs in California."; +//! +//! let entities = gliner.extract(text, &labels).await?; +//! for entity in entities { +//! println!("{}: {} ({:.2})", entity.label, entity.text, entity.score); +//! } +//! Ok(()) +//! } +//! ``` +//! +//! ## Label Caching +//! +//! For production workloads with fixed label sets, you can pre-compute label +//! embeddings for significant speedup: +//! +//! ```rust, no_run +//! use rgliner::*; +//! +//! # async fn example() -> anyhow::Result<()> { +//! let mut gliner = Gliner::new().await?; +//! +//! // Pre-compute label embeddings once +//! let labels = ["person", "organization", "location"]; +//! gliner.cache_labels(&labels).await?; +//! +//! // Fast inference with cached labels +//! let documents = ["Apple Inc. was founded by Steve Jobs.", "Microsoft is in Seattle."]; +//! for text in documents { +//! let entities = gliner.extract_with_cached_labels(text).await?; +//! // Process entities... +//! } +//! # Ok(()) +//! # } +//! ``` +//! +//! ## Relation Extraction (GLiNER-RelEx) +//! +//! For joint NER and relation extraction, use the `relex` module: +//! +//! ```rust, no_run +//! use rgliner::relex::*; +//! +//! # async fn example() -> anyhow::Result<()> { +//! let relex = GlinerRelEx::builder() +//! .with_source(GlinerRelExSource::relex_multi()) +//! .build() +//! .await?; +//! +//! let (entities, relations) = relex.extract( +//! "Apple was founded by Steve Jobs.", +//! &["person", "organization"], +//! &["founded by"], +//! ).await?; +//! # Ok(()) +//! # } +//! ``` + +#![warn(missing_docs)] + +mod config; +mod decoding; +mod error; +mod raw; +pub mod relation_decoding; +pub mod relex; +pub mod relex_tokenization; +mod source; +mod tokenization; + +pub use config::GlinerConfig; +pub use decoding::{Decoder, DecodingMode, Entity}; +pub use error::{GlinerError, GlinerLoadingError}; +pub use source::GlinerSource; + +/// Deduplicate entities appearing in more than one overlapping chunk, keeping the +/// highest-scoring occurrence and sorting by span position. +pub(crate) fn merge_entities(entities: &mut Vec) { + entities.sort_by(|a, b| { + a.start_char + .cmp(&b.start_char) + .then_with(|| a.end_char.cmp(&b.end_char)) + .then_with(|| a.label.cmp(&b.label)) + }); + entities.dedup_by(|b, a| { + if a.start_char == b.start_char && a.end_char == b.end_char && a.label == b.label { + if b.score > a.score { + a.score = b.score; + } + true + } else { + false + } + }); +} + +use fusor::{Device, Tensor, VarBuilder}; +use kalosm_common::Cache; +use kalosm_model_types::{FileSource, ModelLoadingProgress}; +use rbert::BertSource; +use std::sync::Arc; +use tokenizers::Tokenizer; + +use crate::raw::{CachedLabels, LabelEncoder, SpanLayer, TextEncoder}; +use crate::tokenization::{first_subtoken_pooling, TokenizedText, WordTokenizer}; + +pub(crate) async fn default_device() -> Device { + Device::gpu().await.unwrap_or_else(|_| Device::cpu()) +} + +/// Download a model artifact from `source` through `cache`, reporting progress +/// under the label `"