Skip to main content

TextEncoder

Trait TextEncoder 

Source
pub trait TextEncoder: Sync {
    // Required methods
    fn n_embd(&self) -> usize;
    fn n_ctx_train(&self) -> usize;
    fn pooling_type(&self) -> PoolingType;
    fn encode_on_worker(
        &self,
        tokens: &[u32],
        segments: Option<&[u32]>,
    ) -> Result<Vec<f32>, EncodeError>;

    // Provided methods
    fn wrap_special(&self, pieces: &[u32]) -> Vec<u32> { ... }
    fn n_segments(&self) -> usize { ... }
    fn wrap_special_pair(&self, _a: &[u32], _b: &[u32]) -> Option<PairSequence> { ... }
    fn encode(
        &self,
        tokens: &[u32],
        segments: Option<&[u32]>,
    ) -> Result<Vec<f32>, EncodeError> { ... }
    fn encode_tokens(&self, tokens: &[u32]) -> Result<Vec<f32>, EncodeError> { ... }
    fn embed_tokens(&self, tokens: &[u32]) -> Result<Vec<f32>, EncodeError> { ... }
}
Expand description

A model that turns a whole token sequence into hidden states in one pass, with no carried state and no logits.

Sync because TextEncoder::encode hands &self to a rayon worker for the duration of the pass; see its doc comment.

Required Methods§

Source

fn n_embd(&self) -> usize

Width of one hidden-state row, and of the pooled embedding.

Source

fn n_ctx_train(&self) -> usize

Longest sequence this checkpoint can represent. For a learned position table this is the table’s height, and exceeding it is an error rather than a degradation.

Source

fn pooling_type(&self) -> PoolingType

What the checkpoint’s own {arch}.pooling_type said.

Source

fn encode_on_worker( &self, tokens: &[u32], segments: Option<&[u32]>, ) -> Result<Vec<f32>, EncodeError>

n_tokens × n_embd hidden states, in row order.

segments is the per-position segment id, or None for “all zeros” — the single-sequence case. There is deliberately ONE graph with a parameter rather than a segment-aware copy of a segment-blind one: this repo has lost a model feature to every copied forward pass it has ever had, and the pair path is exercised far less often than the embedding path, so a copy is exactly where a fix would fail to land.

This is the body an encoder writes; callers want Self::encode, which is the same computation with the CPU worker pool entered once for the whole pass.

Provided Methods§

Source

fn wrap_special(&self, pieces: &[u32]) -> Vec<u32>

Wraps a tokenizer’s pieces in whatever the model requires around them — for BERT, [CLS] … [SEP].

This is on the encoder rather than the tokenizer because ferrox’s tokenizers deliberately encode text only (they are checked token-for-token against llama_tokenize(..., add_special = false, ...)), while llama.cpp keeps add_special in the vocab and applies it here. The default adds nothing, so an encoder that genuinely needs no wrapper does not have to say so.

Source

fn n_segments(&self) -> usize

How many rows the checkpoint’s token-type (“segment”) table carries, i.e. the number of distinct segment ids Self::encode will accept.

1 — the default — means the encoder can only represent “Sentence A”, which is enough for an embedding pass and is not enough for a cross-encoder pair. Read at load time by crate::EmbeddingModel, which refuses a reranker checkpoint that cannot express segment 1 rather than silently scoring the document half as segment 0.

Source

fn wrap_special_pair(&self, _a: &[u32], _b: &[u32]) -> Option<PairSequence>

The two-segment input a cross-encoder scores: one sequence holding a query and a document with the model’s own boundary between them, and the segment id of every position. For BERT that is [CLS] a [SEP] b [SEP] with segments 0…0 1…1, which is exactly what HuggingFace’s tokenizer(query, document) emits.

None — the default — means this encoder has no two-segment form, and a caller that needs one must refuse. Deliberately NOT defaulted to wrap_special(a ++ b): a cross-encoder was trained with a separator between the halves, and one that never sees it still returns a plausible float. That is the “computes something else” failure, and it is invisible — a rerank with no boundary produces an ordering, just not the model’s.

Source

fn encode( &self, tokens: &[u32], segments: Option<&[u32]>, ) -> Result<Vec<f32>, EncodeError>

Self::encode_on_worker, with the CPU worker pool entered once for the whole pass.

Same rule, and the same reason, as crate::engine::Engine::forward_token: a forward pass through a stack of quantized projections opens a parallel region per matmul, and every one of them costs a pthread park and wake when the driving thread is not a rayon worker. An encoder pass is one pass over the whole sequence rather than one per token, so the saving is smaller than a decode loop’s – but it is the same saving, and an encoder that had to remember to ask for it would be one more place for the rule to be forgotten.

Do not override. See ferrox_core::par::on_workers for the three cases it declines to promote.

Source

fn encode_tokens(&self, tokens: &[u32]) -> Result<Vec<f32>, EncodeError>

Self::encode for a single sequence: every position is segment 0.

Source

fn embed_tokens(&self, tokens: &[u32]) -> Result<Vec<f32>, EncodeError>

Self::encode_tokens followed by the checkpoint’s own pooling. Not L2-normalized — see crate::pooling::pool.

Dyn Compatibility§

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§