README.md

metadata

license: gemma
base_model: google/gemma-4-E2B-it
pipeline_tag: text-generation
tags:
  - gemma4
  - hybrid
  - custom_code
  - custom_generate

gemma-4-e2b-it-hybrid

Cactus-Compute/gemma-4-e2b-it-hybrid is google/gemma-4-E2B-it plus a handoff probe: a small head (weight prefix handoff_probe.*) that scores every generation with

confidence = 1 - p_wrong

so a client can decide when to keep the on-device answer and when to hand off to a cloud model. The base weights are byte-identical to the stock checkpoint (same keys); the repo adds eleven probe tensors, a remote-code model class and a custom_generate recipe. Stock engine commands work unchanged — you only add --trust-remote-code / trust_remote_code=True.

Serving with stock transformers

transformers serve --trust-remote-code
# then request model "Cactus-Compute/gemma-4-e2b-it-hybrid" via the OpenAI-compatible API

or interactively:

transformers chat Cactus-Compute/gemma-4-e2b-it-hybrid --trust-remote-code

How the confidence reaches you (in-band trailer)

transformers serve cannot add response fields, so the score travels in-band: the assistant content's final line is

\n[[hybrid:confidence=0.7812]]

always exactly 4 decimals, ASCII, confidence in [0, 1]. Strip the final [[hybrid:...]] line before display and parse the float for routing. The trailer is emitted in both streaming and non-streaming modes.

Limitations

• --continuous-batching: not supported — the CB scheduler bypasses generate(), so no probe runs and no trailer is emitted. Serve without --continuous-batching to get confidence scores.
• The trailer (and confidence) is only produced for single-sequence decoding: batch size 1, no beam search, no assisted/speculative decoding. Unsupported modes fall back to stock behavior (no trailer).
• The probe scores at most the first 1024 generated tokens.
• The trailer's token ids are appended to the returned sequences, so reported completion token counts include the trailer (a handful of tokens).

Python usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "Cactus-Compute/gemma-4-e2b-it-hybrid", trust_remote_code=True, dtype="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Cactus-Compute/gemma-4-e2b-it-hybrid")
inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": "Explain why this contract clause is risky."}],
    add_generation_prompt=True, return_tensors="pt",
).to(model.device)

# Structured API: clean sequences + raw float (no in-band trailer).
sequences, confidence = model.generate_with_confidence(inputs, max_new_tokens=512)
print(confidence)                     # e.g. 0.7812
print(model.last_confidence)          # same value

# Stock generate (custom_generate recipe): plain tensor + in-band trailer.
sequences = model.generate(inputs, max_new_tokens=512)

# Or opt into a rich return with the raw float:
out = model.generate(inputs, max_new_tokens=512, return_confidence=True)
out.sequences, out.confidence, out.trailer_text

# Suppress the trailer while keeping stock behavior:
sequences = model.generate(inputs, max_new_tokens=512, emit_trailer=False)

Probe contract

• Input: float32 [T, 1536] — output of decoder layer index 28 (config.probe_layer), captured at the position that predicts each generated token: row 0 = last prompt position at prefill, row t = position captured at generation step t. Only the first 1024 rows are scored.
• Math (float32): x = LayerNorm(x, eps=1e-5) * norm.weight + norm.bias; p = relu(x @ proj.weight.T + proj.bias); s = p @ attn_query / sqrt(32); w = softmax_T(s - max); pooled = w @ p; h = relu(head.0 @ pooled + b); h = relu(head.2 @ h + b); logit = head.4 @ h + b; p_wrong = sigmoid(logit); confidence = 1 - p_wrong.
• Capture uses a forward hook that keeps only one [1, 1536] row per decode step — full hidden-state stacks are never materialized.

MLX (Apple Silicon)

The repo ships a single-file mlx-lm model (gemma_4_e2b_it_hybrid.py), referenced by "model_file" in config.json (mlx-lm >= 0.30.1). After generation, read model.last_confidence (or model.confidence(num_tokens=N)). An mlx-lm Model cannot inject tokens into the stream, so there is no in-band trailer on MLX today — confidence is Python-API only there.

Repo contents

File	Purpose
configuration_gemma_4_e2b_it_hybrid.py	Gemma4E2BItHybridConfig (stock Gemma-4 text config + probe hyperparams)
modeling_gemma_4_e2b_it_hybrid.py	Gemma4E2BItHybridForCausalLM (stock Gemma4ForCausalLM + handoff_probe.*)
custom_generate/generate.py	stock decode loop + confidence + in-band trailer
model*.safetensors	base weights (identical keys) + handoff_probe.* tensors
gemma_4_e2b_it_hybrid.py	single-file mlx-lm model, wired via config.json's model_file

Requires transformers>=5,<6. Gemma model use remains subject to the Gemma terms.