Text Generation, Embeddings, Rerankers, and Classifiers

“Call the AI model” is not an interface.

A text generator returns one token decision at a time and may run for hundreds of decode steps. An embedding model returns one fixed-width vector. A reranker scores each query/candidate pair. A classifier returns scores over a declared label space. Those outputs have different shapes, semantics, latency curves, batching strategies, failure modes and evaluation metrics.

They may share Transformer layers, tokenizers or model families. That does not make them interchangeable.

This chapter adds checked output contracts to mosaic-ml and runs all four against deterministic fixture outputs. The code starts at the application boundary after a runtime has produced logits, vectors or scores. It does not pretend that a few arrays are a trained neural network. The purpose is to make every adapter preserve the task semantics that a real local or remote runtime must supply.

Learning objectives

By the end you will be able to:

• state the input/output shape of generation, embedding, reranking and classification;
• distinguish a backbone from its task-specific head and training objective;
• map one vocabulary-sized generation logit vector to one token decision;
• explain why generation is a repeated stateful operation rather than one classification call;
• normalize finite embeddings and calculate cosine similarity;
• keep embedding model, dimensions, normalization and query/document roles in index metadata;
• explain bi-encoder retrieval and cross-encoder reranking;
• preserve raw reranker scores without calling them probabilities;
• convert mutually exclusive class logits with stable softmax;
• distinguish single-label softmax from multi-label sigmoid;
• retain label maps as model-release artifacts;
• design stable tie-breaking and bounded top_k;
• reject non-finite values, mismatched shapes, duplicates and excessive collections;
• choose a specialized task model instead of using generation for every problem;
• compose retrieve → rerank → decide → generate without erasing provenance;
• select task-appropriate quality, latency and capacity metrics.

Completion artifact

Run:

cd rust/rust-ai-engineering
cargo run -q -p mosaic-ml --example task_output_contracts

The example prints four different contracts:

{
  "generation": {
    "contract": "one vocabulary-sized logit vector -> one token decision",
    "token_id": 2,
    "piece": " safe",
    "is_eos": false
  },
  "embedding": {
    "contract": "one input -> one fixed-width vector",
    "dimensions": 3,
    "cosine_similarity": 1.0
  },
  "reranking": {
    "contract": "query + each candidate -> one raw relevance score",
    "top": [
      { "candidate_id": "doc-leases", "rank": 1, "raw_score": 2.4 },
      { "candidate_id": "doc-retries", "rank": 2, "raw_score": 1.7 }
    ]
  },
  "classification": {
    "contract": "one input -> one mutually-exclusive label distribution",
    "labels": [
      { "label": "answer", "probability": 0.8017, "rank": 1 },
      { "label": "needs_tool", "probability": 0.1325, "rank": 2 },
      { "label": "abstain", "probability": 0.0658, "rank": 3 }
    ]
  }
}

The shortened decimals above are for reading. The real JSON preserves the calculated floating-point values. If a wire contract needs fixed precision, define its rounding and integer representation explicitly as the calibration chapter did.

One table before any abstraction

Task	Logical input	Raw model output	Application output	Repeated?
Generation	tokenized prefix + state	one logit per vocabulary token	selected next token, then text/structure	once per generated token
Embedding	one query/document/item	fixed-width dense vector	stored or compared representation	once per item
Reranking	query + one candidate	scalar/logit per pair	ordered candidate IDs	once per candidate
Single-label classification	one item	one logit per label	one normalized label distribution	once per item

The table already exposes four common bugs:

1. using a classifier probability as if it were a generation token probability;
2. inserting vectors with a new dimension into an old index;
3. applying a reranker to every document rather than a retrieved candidate set;
4. applying softmax to a multi-label problem whose labels are not mutually exclusive.

Backbone, objective, pooling, and head

A Transformer backbone produces contextual hidden states. The deployable task behavior also depends on what happens around those states:

input serialization
  → tokenizer / position handling
  → Transformer backbone
  → pooling or task head
  → activation / decoding
  → application contract

Examples:

• a causal language-model head projects each relevant hidden state to vocabulary logits;
• an embedding model pools token states, may apply a projection, and is trained so useful pairs have suitable geometry;
• a cross-encoder reranker jointly encodes a query/candidate pair and projects to one or more relevance logits;
• a sequence classifier pools a sequence representation and projects to declared labels.

The same parameter count or base architecture does not imply the same behavior. Training loss and data determine what distances or scores mean. A random mean pool of a general language model’s hidden states is not automatically a high-quality search embedding model.

The task head, pooling rules, label map, input prefixes and normalization policy belong in the model release contract from the model-artifact chapter.

Contract 1: autoregressive text generation

At one decode step, the language-model head produces:

logits shape = [vocabulary_size]

For a batch, a runtime may expose [batch, sequence, vocabulary] or only the final-position logits needed for selection. The application must know which position and batch row it received.

The companion validates a display vocabulary:

pub struct TokenVocabulary {
    pieces: Vec<String>,
    eos_token_id: usize,
}

Then:

let step = select_generation_token(
    &vocabulary,
    &[0.1, 0.2, 1.4, -0.5],
    &mut Sampler::new(42),
    SelectionStrategy::Greedy,
)?;

The function first proves:

logits.len() == vocabulary.len()

It then delegates selection to the versioned sampler from the earlier chapter and maps the selected ID back to a token piece. It also compares the ID with the configured EOS ID.

This is one step, not “generate text”:

while not stopped:
    run next-position forward pass
    process logits
    select one token
    append token and KV state
    update stop detectors
    stream safe output

A complete generator needs:

• maximum-new-token and total-context enforcement;
• EOS and possibly multiple stop-token handling;
• string stop detection across token boundaries;
• cancellation and deadline checks;
• incremental decoding without corrupting partial UTF-8;
• structured-output constraints or validators;
• tool-call parsing;
• usage accounting;
• terminal reason: EOS, length, stop sequence, cancellation, policy or failure.

Token pieces are not identifiers

The first focused example gate found a real modeling error. The initial validator reused the human-label rule name.trim() == name for vocabulary pieces. It rejected " is" and " safe". Leading spaces are meaningful in many tokenizers.

The corrected boundary has separate validation:

// Human labels/document IDs:
non-empty + bounded + no control characters + no surrounding whitespace

// Teaching vocabulary pieces:
non-empty + bounded + no control characters
// surrounding spaces are preserved

Even this display vocabulary is narrower than a production tokenizer. Byte-fallback tokenizers can represent pieces that are not independently displayable Unicode. The tokenizer/decoder must own the exact byte-to-text contract; do not build it from a sanitized logging label.

Generation score semantics

Logits are unnormalized scores. Softmax at one step gives a distribution over the next token under the rendered prefix and model. It is not:

• the probability the final answer is correct;
• the probability a factual claim is true;
• directly comparable across different prefixes;
• a calibrated task-success estimate.

Generation quality also depends on the sampler, stop policy and later validation. Evaluate the complete trajectory.

Cost shape

Generation has prefill plus repeated decode. Output length changes latency and cost. Capacity needs KV-cache space per active sequence. Report TTFT, inter-token latency, output throughput, stop reason and token counts—not just requests per second.

Contract 2: embeddings

An embedding maps one item to a fixed-width dense vector:

embed(input) -> [d]

Common uses include semantic retrieval, clustering, duplicate detection, recommendation and inputs to another model. Geometry only has meaning under the model’s training and post-processing contract.

The companion accepts finite f32 components and produces an L2-normalized f64 vector:

let query = UnitEmbedding::try_from_f32(vec![3.0, 4.0, 0.0])?;
let document = UnitEmbedding::try_from_f32(vec![6.0, 8.0, 0.0])?;
let similarity = query.cosine_similarity(&document)?;

Normalization computes:

norm(x) = sqrt(Σ x_i²)
unit(x) = x / norm(x)

For unit vectors:

cosine(x, y) = unit(x) · unit(y)

The example vectors point in the same direction, so cosine similarity is 1.0.

The constructor rejects:

• an empty vector;
• more than 65,536 dimensions in this teaching boundary;
• NaN or infinity;
• zero or effectively zero norm;
• non-finite normalization results.

Similarity rejects dimension mismatch and clamps tiny floating-point overshoot to [-1, 1].

Index identity is more than dimension

Store alongside every vector/index:

embedding model immutable revision
tokenizer/processor revision
pooling and projection
dimension
normalization policy
distance metric
query/document role or prefix
input truncation policy
quantization/storage format
created-at and corpus snapshot

Two models can both emit 768 values and inhabit incompatible spaces. Silent mixing may produce numerically valid, semantically useless results.

Changing an embedding release usually requires shadow indexing or re-embedding the corpus. A safe rollout builds a new versioned index, evaluates it, switches reads, and retains rollback until the old index can be retired.

Query and document roles

Some retrieval models are symmetric: the same transformation is appropriate for both sides. Asymmetric retrieval models may require distinct prompts, prefixes or routes for queries and documents. Treat embed_query and embed_document as different capabilities even if they return the same dimension.

Accidentally embedding documents with the query prefix does not create a type or shape error. It creates a quality failure. Include role-specific golden cases.

Similarity is not relevance probability

A cosine value of 0.82 is not automatically “82% relevant.” Score distributions vary with model, domain, corpus, query and preprocessing. Tune thresholds on held-out data and measure retrieval metrics at the operating point.

Useful metrics include:

• recall@k: was a relevant candidate retrieved?
• precision@k;
• mean reciprocal rank;
• normalized discounted cumulative gain;
• latency, items/second and index memory;
• failure rates by domain/language/modality.

Contract 3: reranking

A reranker evaluates a query and candidate together:

score(query, candidate_i) -> scalar_i

A cross-encoder can jointly attend across both texts, which often makes its relevance decision stronger than a bi-encoder similarity. The cost is that it must run for every pair.

That leads to the standard two-stage shape:

query
  → inexpensive lexical/vector retrieval over large corpus
  → top N candidates
  → cross-encoder score for each of N pairs
  → stable top K

Do not rerank a million documents one pair at a time. First-stage retrieval supplies coverage; reranking improves ordering inside a bounded candidate set.

The companion begins after pair scoring:

let ranked = rerank(
    vec![
        CandidateScore { candidate_id: "doc-retries".into(), score: 1.7 },
        CandidateScore { candidate_id: "doc-borrowing".into(), score: -0.2 },
        CandidateScore { candidate_id: "doc-leases".into(), score: 2.4 },
    ],
    2,
)?;

It rejects empty/oversized sets, invalid IDs, duplicates, non-finite scores, zero top_k, and top_k above the candidate count. It sorts descending by raw score, then ascending by stable candidate ID for a deterministic tie. It retains original index and emits explicit rank.

Do not relabel logits as probabilities

Some rerankers emit logits. Applying sigmoid maps a scalar to (0, 1) but does not prove empirical calibration. It also does not change ordering because sigmoid is monotonic. The companion therefore names the field score, preserves negative values, and does not normalize scores across candidates.

Softmax across a candidate list would make a document’s reported value depend on which other documents happened to be present. That may be useful for a particular listwise model, but it is not a generic conversion from pairwise score to relevance probability.

Candidate-set coupling

Reranker quality cannot repair a missing relevant document. Measure:

first-stage recall@N
reranked nDCG@K / MRR@K
end-to-end answer or task quality
pair count and token lengths
reranker batch latency

Increasing N can improve recall but grows pairwise cost and may introduce hard distractors. Tune retrieval depth and rerank depth together.

Preserve source IDs through ranking. The generator should cite the exact selected documents, not freshly search after answering.

Contract 4: single-label classification

Single-label classification chooses exactly one label from a fixed set:

logits shape = [label_count]
labels       = [needs_tool, answer, abstain]

The companion uses numerically stable softmax:

m = max(logits)
p_i = exp(logit_i - m) / Σ exp(logit_j - m)

Subtracting the maximum prevents overflow without changing the distribution. Output predictions are sorted by probability and stable label name, truncated to validated top_k, and assigned ranks.

Label order is part of the model release. If a head was trained with:

0 = abstain
1 = answer
2 = needs_tool

but the application loads a different mapping, all shapes still match while every meaning is wrong. Digest the mapping with the model artifact.

The probability is a softmax value, not automatically a calibrated chance of correctness. Measure calibration on held-out cases and use the decision policy from the previous chapter.

Single-label versus multi-label

The companion is deliberately single-label. Softmax couples labels so their probabilities sum to one.

Multi-label tasks allow several independent labels:

["contains_code", "security_sensitive", "requires_network"]

They commonly apply sigmoid independently to each label, then use per-label thresholds:

sigmoid(z_i) = 1 / (1 + exp(-z_i))

Do not use the companion’s classify_single_label for multi-label output. Add a separately named contract with label-specific thresholds and metrics.

Classification metrics

Accuracy is appropriate only when classes and costs make it informative. Also consider:

• per-label precision and recall;
• macro versus micro F1;
• confusion matrix;
• calibration/Brier for declared events;
• false-action cost;
• abstention/coverage;
• latency, batch throughput and memory;
• shifted and adversarial slices.

For a router, false answer on a tool-required case may cost far more than false needs_tool. Release policy should reflect the cost matrix.

Why not generate a label as text?

A general generator can be prompted to emit "answer", "abstain" or "needs_tool". That can be useful when labels change frequently, demonstrations carry rich semantics, or one generative call must also explain/structure a decision.

But a specialized classifier often offers:

• smaller artifacts;
• one forward pass rather than autoregressive decoding;
• easy batching;
• fixed label mapping;
• direct logit access;
• lower cost and latency;
• clearer confusion-matrix and calibration analysis.

Generation adds failure modes:

• extra prose around the label;
• spelling/casing variants;
• invalid schema;
• prompt injection;
• sampling variance;
• stop failures;
• more compute.

Choose from evidence. Benchmark a constrained generator, classifier and hybrid router on the same held-out cases. Small models may classify or rerank extremely well even when they are weak open-ended generators.

A typed pipeline without a universal score

Consider an evidence-grounded assistant:

1. classifier
   decide answer / tool / abstain

2. query embedder
   map the search query into the versioned retrieval space

3. first-stage index
   retrieve candidate IDs with source metadata

4. reranker
   score query/candidate pairs and select top evidence

5. generator
   produce a cited structured answer from selected evidence

6. validators
   verify schema, citations, policy, grounded claims and budgets

Never place a single score: f32 on the shared response and hope callers infer meaning. Prefer distinct types:

GenerationStep       { token_id, piece, stop state, sampler evidence }
UnitEmbedding        { values, release/dimension contract }
RerankedCandidate    { source ID, raw score, original index, rank }
ClassPrediction      { label, probability, rank }

The provider-neutral traits later in this part will preserve these capabilities rather than expose a god-object method.

Bounds, security, and deterministic ordering

Model outputs cross a trust boundary. Even local files can be corrupt; remote providers can change; an adapter can deserialize malicious sizes.

This module imposes teaching limits:

vocabulary        ≤ 262,144 pieces
embedding         ≤ 65,536 dimensions
classification    ≤ 16,384 labels
rerank set        ≤ 65,536 candidates
name/token bytes  ≤ 256

Production limits should come from the pinned model contract and service capacity, usually much tighter per endpoint. Check lengths before allocating transformed output.

Every raw number must be finite. NaN breaks equality and ordering; infinity can dominate softmax or ranking. Failing closed is safer than silently replacing invalid values.

Stable tie-breaking matters because:

• equal quantized scores are common;
• unstable order creates flaky goldens and pagination;
• downstream caching/replay needs a canonical result;
• rank changes can alter which evidence reaches generation.

The chosen secondary key is label or candidate ID. Record the policy as part of the adapter version.

Do not log raw text just because a score fails. Log bounded identifiers, shape, release and safe error code; handle sensitive content through the artifact/redaction policy.

Performance planning by task

Generation

Batch prefill and continuously batch decode where the runtime supports it. Plan KV cache, output length and streaming. Track TTFT and inter-token latency.

Embedding

Batch many similarly sized inputs, bucket by length, avoid padding waste, and stream vectors into a versioned index. Throughput is items or tokens per second; storage grows roughly with items × dimensions × bytes/component, plus index overhead.

Reranking

Batch query/candidate pairs while bounding candidates × pair_tokens. Cache only when query, candidate content digest, model release and preprocessing are identical and policy permits it.

Classification

Batch fixed-head inference. A compact encoder can handle much more traffic than a generator. Tune the decision threshold/cost policy separately from runtime batching.

Device utilization is not user value. Measure end-to-end retrieval/decision/answer quality and latency under the real length distribution.

Failure investigation

Symptom	First checks
Generated text has missing spaces	tokenizer decoder and preservation of boundary-space/byte pieces
Retrieval quality collapses after release	model/index revision, dimensions, query/document role, normalization
Reranker values look like probabilities but exceed 1	raw logit semantics and accidental field naming
Relevant document never reaches top K	first-stage recall before blaming reranker
Classification labels appear permuted	immutable label-ID mapping
Softmax returns NaN	non-finite logits, unstable exponentiation, corrupted runtime output
Results reorder between runs on ties	missing canonical secondary key
Latency spikes with corpus size	reranker accidentally applied before bounded retrieval

Investigate the earliest violated invariant. A later generator cannot restore a candidate that retrieval discarded.

Exercises

Exercise 1: identify the task

For each output, name the contract:

1. [0.12, -0.30, 1.44, ...] with 1,024 fixed entries for every document.
2. one scalar for each (query, document) pair.
3. a 12-element vector whose positions map to intent labels.
4. 150 repeated vocabulary distributions for one response.

Solution

1 is an embedding; 2 is reranking; 3 is classification; 4 is autoregressive generation. The raw arrays are insufficient without model revision, mapping, normalization and preprocessing metadata.

Exercise 2: build multi-label classification

Add classify_multi_label. Use stable sigmoid, one threshold per immutable label, finite-value checks and stable label ordering. Test a case where zero, one and multiple labels pass.

Design guidance

Do not normalize across labels. Validate threshold count equals label count and each threshold is in [0,1]. Return each label’s independent score plus whether it passed. Evaluate per-label precision, recall and calibration.

Exercise 3: version an embedding index

Define a Rust EmbeddingIndexContract containing artifact digest, dimension, metric, normalization, query/document preprocessing revision and corpus snapshot. Reject insertion/search on mismatch.

Exercise 4: retrieve then rerank

Create ten candidate embeddings, retrieve top five by cosine, rerank those five with supplied pairwise scores, and return top two. Prove in a test that the reranker never sees the other five.

Exercise 5: preserve byte-level tokens

Replace TokenVocabulary’s display String pieces with Vec<u8> pieces and an incremental UTF-8 decoder. Test a Unicode scalar split across multiple tokens and invalid terminal bytes.

Exercise 6: compare cost

For 500 requests/second, compare:

• classifier: one 8 ms batchable pass;
• generator: 10 ms prefill plus 12 tokens at 4 ms/token;
• retrieve/rerank: 4 ms retrieval plus 50 pairs at an amortized 0.4 ms/pair.

State which operations can batch, the concurrency implied by Little’s Law, and which quality metric must accompany the latency result.

Check your understanding

• Why can two equal-dimension embedding models still be incompatible?
• Why does sigmoid not automatically turn a reranker logit into a calibrated relevance probability?
• What changes between single-label and multi-label classification?
• Why is text generation stateful across output tokens?
• Which metric catches a relevant document omitted before reranking?
• Why are leading spaces valid token pieces but suspicious document IDs?
• What should break a cache key for a reranker result?
• Why does one universal score field destroy semantics?

Completion checklist

• generation logits exactly match vocabulary size;
• EOS and every stop reason are explicit;
• tokenizer pieces preserve their byte/text semantics;
• embedding vectors are finite, non-zero and dimension checked;
• model/index identity includes role, metric and normalization;
• first-stage retrieval and reranking remain separate stages;
• reranker scores remain raw unless a calibrated mapping is proven;
• label-ID maps are immutable release artifacts;
• single-label softmax is not reused for multi-label tasks;
• all top_k values are bounded and non-zero;
• duplicate IDs and non-finite values fail closed;
• ties have a canonical secondary key;
• each task has quality, latency, throughput and memory metrics;
• specialized models are compared with a generative baseline on held-out data.

Sources and further study

• Hugging Face Transformers, Generation: generation strategies, logits processing, stopping criteria and cache options.
• Hugging Face Transformers, Pipelines: task-specific pipeline inputs/outputs across generation, feature extraction and classification.
• Devlin et al., BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding: a bidirectional encoder adapted to downstream tasks with task output layers.
• Reimers and Gurevych, Sentence-BERT: bi-encoder sentence representations for efficient similarity search.
• Sentence Transformers, Semantic Search: query/document encoders, asymmetric retrieval and similarity search.
• Sentence Transformers, Cross-Encoder Usage: pairwise scoring, reranking cost, logits and optional activations.
• Sentence Transformers, Retrieve & Re-Rank: the bi-encoder retrieval plus cross-encoder reranking architecture.

What this chapter does not claim

The fixture arrays are not learned model outputs and the code is not an optimized inference engine, vector database or tokenizer. L2 normalization is not appropriate for every embedding model. Softmax probabilities are not calibrated automatically. Raw reranker scores do not transfer between models or queries. The chapter proves something narrower and essential: four task families require four honest contracts, checked shapes and task-specific evaluation.

The next chapter pins model metadata, immutable revisions and provider changes so these contracts cannot silently drift underneath a stable application name.