undercurrent.model

The front door. ProbedModel loads a model, attaches the probes your spec asks for, and returns GenerationOutputs with each extraction point's result. ProbedModel and GenerationOutput are also exported from undercurrent.

from undercurrent import ProbedModel

Loading and generating

ProbedModel(model=None, *, spec=None, probes=None, backend='hf', device=None, on_result=None, log_sink=None, metrics_sink=None, router_kwargs=None, router=None, max_concurrency=None, **model_kwargs)

A model with probes attached. Build it with ProbedModel.from_pretrained(...) (the constructor takes the same arguments).

from undercurrent import ProbedModel

with ProbedModel.from_pretrained("openai-community/gpt2", spec="probes.yaml") as model:
    out = model.generate("The weather today is", max_new_tokens=20)
print(out.text, out.probe_results)

Use it as a context manager, or call close() when done.

Parameters:

Name	Type	Description	Default
model	Any	a Hugging Face hub id or local path, or an already-built model object (then pass tokenizer=). May be None only when backend is an EngineAdapter that already has a model loaded.	None
spec	SpecInput	the extraction points: a YAML file path, a YAML string, a dict, a ProbeSpec, a list of ExtractionPoint, or None for no probes.	None
probes	Mapping[str, type[Probe] | ProbeFactory] | ProbeRegistry | None	{probe_type: Probe subclass | ProbeFactory}. Defaults to the global registry (@register_probe). Can't be combined with router=.	None
backend	str | EngineAdapter	"hf" (transformers, the default), "vllm" (needs vLLM installed in the environment, in a supported version; vLLM places the model itself and **model_kwargs are vLLM engine args), or an EngineAdapter instance (used as-is; the caller keeps owning it).	'hf'
device	str | None	device for the model, e.g. "cpu" or "cuda" (HF default: cpu). vLLM places the model itself; leave it None.	None
max_concurrency	int | None	how many prompts of a batch generate() runs at once. Defaults to the backend's own value (HF: 1, vLLM: DEFAULT_VLLM_MAX_CONCURRENCY); HF only supports 1.	None
on_result	OnResult | None	callback(GenerationOutput), called once per finished generation on the thread that finished it. Exceptions it raises are logged on the undercurrent logger and never break generation.	None
log_sink	SupportsLogSink | None	advanced; a LogSink (e.g. a FileLogSink) for async ("observe mode") extraction points.	None
metrics_sink	MetricsSink | None	advanced; a MetricsSink for async-binding metrics.	None
router_kwargs	Mapping[str, Any] | None	advanced; forwarded to Router (worker_pool_size, default_queue_depth, default_overflow_policy, drain_timeout, default_intervention_policy, circuit_breaker_threshold).	None
router	Router | None	advanced; use this Router as-is. The caller owns it, so close() doesn't shut it down. Can't be combined with probes= or router_kwargs=. log_sink/metrics_sink are attached to it.	None
**model_kwargs	Any	forwarded to the backend's model loading (for HF, tokenizer= plus from_pretrained kwargs such as torch_dtype; for vLLM, engine args such as gpu_memory_utilization and max_model_len).	{}

Raises:

Type	Description
ProbedModelConfigError	a probe_type doesn't resolve, a probe_kind doesn't match its probe, a tensor_type isn't supported by the backend, or a layer index is out of range.
MissingDependencyError	backend="vllm" but vLLM isn't installed.

router property

The Router dispatching activations to probes (advanced use).

adapter property

The engine adapter running the model (advanced use).

spec property

The resolved extraction points every generation is probed with.

max_concurrency property

How many prompts of a batch generate() runs at once (1 for HF).

from_pretrained(model=None, **kwargs) classmethod

Load model and attach probes; see the class docstring for the arguments.

generate(prompt, *, max_new_tokens=None, temperature=None, top_p=None, seed=None, stop=None, return_exceptions=False, **backend_kwargs)

generate(prompt: str, *, max_new_tokens: int | None = ..., temperature: float | None = ..., top_p: float | None = ..., seed: int | None = ..., stop: str | Sequence[str] | None = ..., return_exceptions: bool = ..., **backend_kwargs: Any) -> GenerationOutput

generate(prompt: Sequence[str], *, max_new_tokens: int | None = ..., temperature: float | None = ..., top_p: float | None = ..., seed: int | None = ..., stop: str | Sequence[str] | None = ..., return_exceptions: Literal[False] = ..., **backend_kwargs: Any) -> list[GenerationOutput]

generate(prompt: Sequence[str], *, max_new_tokens: int | None = ..., temperature: float | None = ..., top_p: float | None = ..., seed: int | None = ..., stop: str | Sequence[str] | None = ..., return_exceptions: Literal[True], **backend_kwargs: Any) -> list[GenerationOutput | Exception]

Generate for one prompt (returns a GenerationOutput) or a list of prompts (returns a list, same order).

A list runs up to max_concurrency prompts at once (vLLM batches them; HF runs them one by one). By default one failing prompt raises its exception for the whole call (prompts not started yet are cancelled). With return_exceptions=True the exception is put in that prompt's slot of the returned list instead. A single string prompt always raises.

temperature=0 is greedy decoding, temperature > 0 samples. seed makes sampling reproducible. stop ends generation at a string (not included in the text). Other keyword arguments go to the backend unchanged (for HF, model.generate(...) kwargs such as repetition_penalty).

close()

Release the adapter and shut down the Router, unless the caller passed them in (backend=<EngineAdapter>, router=): those stay untouched apart from removing this model's results listener. Idempotent.

GenerationOutput(text, prompt, request_id, probe_results=dict(), aborted=False, abort_reason=None, abort_signal=None, abort_point=None) dataclass

The result of one ProbedModel.generate() call for one prompt.

Attributes:

Name	Type	Description
text	str	the generated text (the prompt not included).
prompt	str	the prompt it was generated from.
request_id	str	the id the probes saw for this generation.
probe_results	Mapping[str, ProbeResult]	{extraction_point_name: ProbeResult} for every extraction point in the spec, async ones included.
aborted	bool	whether an inline probe stopped generation early.
abort_reason	str | None	human-readable reason when aborted, else None.
abort_signal	ProbeSignal | None	the ProbeSignal that stopped generation, if known.
abort_point	str | None	name of the extraction point that stopped generation.

Errors

ProbedModelConfigError

Bases: ProbingError, ValueError

The spec, probes or backend given to ProbedModel don't fit together.

Raised at construction time (fail fast), before any generation. The message names the extraction point and says how to fix it.

Defaults

DEFAULT_MAX_NEW_TOKENS = 32 module-attribute

Used when the generation length isn't given (HF: neither max_new_tokens nor max_length; vLLM: neither max_new_tokens nor max_tokens).

DEFAULT_VLLM_MAX_CONCURRENCY = 64 module-attribute

How many generate() calls the vLLM backend runs at once by default.