interlens.participant.participants¶

def generate_batch(self, views: list[list[dict]], *, turn: int | None = None,
                   group_seed: int | None = None, max_new_tokens: int | None = None) -> list[Message]:
	"""Generate one turn for many independent conversations at once — the API analogue of
	``ModelParticipant.generate_batch``, driven by the runner's co-stepper (``rollout(..., batched=True)``) to
	make large API rollouts cheap and throughput-bound.

	With ``batch=True`` every view is sent as one **asynchronous provider batch** (Anthropic Message Batches /
	OpenAI Batch API) via ``client.submit_batch`` — ~50% cost and much higher throughput, at the price of
	batch-window latency. **If the provider has no batch API (e.g. OpenRouter) this raises** rather than
	silently degrading, so a requested batch is never quietly run as serial calls. With ``batch=False`` it
	falls back to sequential per-view calls (correct, just no batch discount). Interp is unavailable here, as
	for ``generate``. ``turn``/``group_seed`` are accepted for co-stepper compatibility but unused (seeds do
	not bind hosted models). ``metadata['batched']`` marks these turns."""
	if not views:
		return []
	client = self.client or _default_client(self.provider)
	max_tokens = max_new_tokens if max_new_tokens is not None else self.max_tokens
	requests = []
	for view in views:
		system, messages = self._split_view(view)
		requests.append(dict(system=system, messages=messages, model=self.model_id,
		                     max_tokens=max_tokens, temperature=self.temperature))
	if self.batch:
		if not hasattr(client, "submit_batch"):
			raise NotImplementedError(
				f"APIParticipant {self.name!r} has batch=True but its client {type(client).__name__} exposes "
				f"no submit_batch; batch mode is unavailable for provider {self.provider!r}.")
		texts = client.submit_batch(requests)
	else:
		texts = [client(**r) for r in requests]
	return [Message(author=self.name, content=t,
	                metadata={"provider": self.provider, "model": self.model_id, "batched": True})
	        for t in texts]

def parse_tool_calls(self, text: str) -> list:
	"""Parse Gemma's ```` ```tool_code ```` function-call blocks (best-effort).

	Gemma emits calls like ```` ```tool_code\nname(arg="x", n=2)\n``` ```` rather than Hermes JSON. We extract
	the block, then the function name and simple keyword arguments. Falls back to ``[]`` (treat as final
	message) on anything it can't parse, matching the base contract of never misfiring silently.
	"""
	import ast
	import re
	from ...tools.tool_call import ToolCall

	calls = []
	for block in re.findall(r"```tool_code\s*(.*?)```", text, re.DOTALL):
		call_match = re.match(r"\s*([A-Za-z_][\w.]*)\s*\((.*)\)\s*$", block.strip(), re.DOTALL)
		if not call_match:
			continue
		name, arg_src = call_match.group(1), call_match.group(2)
		arguments = {}
		try:
			# Parse ``k=v`` pairs via a synthetic call node, so values are real Python literals.
			parsed = ast.parse(f"f({arg_src})", mode="eval")
			for kw in parsed.body.keywords:
				arguments[kw.arg] = ast.literal_eval(kw.value)
		except (SyntaxError, ValueError) as exc:
			logger.debug("dropping unparseable tool_code block %r: %s", block, exc)
			continue
		calls.append(ToolCall(name=name.split(".")[-1], arguments=arguments, raw=block))
	return calls

def parse_tool_calls(self, text: str) -> list:
	"""Parse Llama-3's ``<|python_tag|>{json}`` function-call format (best-effort).

	Everything after ``<|python_tag|>`` is one or more JSON objects (separated by newlines or semicolons),
	each ``{"name": ..., "arguments"/"parameters": {...}}``. Anything unparseable is skipped, so a malformed
	call yields ``[]`` (treated as a final message), matching the base contract of never misfiring."""
	import json
	from ...tools.tool_call import ToolCall

	tag = "<|python_tag|>"
	idx = text.find(tag)
	if idx < 0:
		return []
	body = text[idx + len(tag):]
	calls = []
	for chunk in body.replace(";", "\n").splitlines():
		chunk = chunk.strip()
		if not chunk:
			continue
		try:
			data = json.loads(chunk)
			name = data["name"]
			arguments = data.get("arguments") or data.get("parameters") or {}
		except (json.JSONDecodeError, KeyError) as exc:
			logger.debug("dropping unparseable python_tag call %r: %s", chunk, exc)
			continue
		calls.append(ToolCall(name=name, arguments=arguments, raw=chunk))
	return calls

def batch_signature(self) -> tuple:
	"""Key identifying which model this participant would batch AS (used by the co-stepper). When loaded, the
	cached model object's identity is authoritative; when not yet loaded, the load recipe is (the cache
	guarantees same recipe on one device -> same object), so participants can be grouped WITHOUT forcing a
	load."""
	if self._model is not None:
		return ("model", id(self._model))
	return ("model", self.weights_path or self.hf_id, str(self.device), self.dtype, self.attn, self.quant,
	        self.revision)

@classmethod
def for_model_type(cls, model_type: str | None) -> type["ModelParticipant"]:
	"""The registered participant class for a transformers ``config.model_type`` (e.g. ``qwen2``, ``gemma2``),
	falling back to base ``ModelParticipant`` for any family that declares no specialized subclass."""
	return cls._REGISTRY.get(model_type or "", ModelParticipant)

@classmethod
def from_model(cls, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase | None = None, *, name: str,
               device: str | torch.device | None = None, **participant_kwargs) -> Self:
	"""Build a participant of THIS class from an already-loaded ``model`` (the eager path — e.g. to wrap a model
	you already hold, or share weights between speakers). ``tokenizer`` is optional — when omitted it is inferred
	from ``model.config._name_or_path``. The chat-template flags (``supports_system_role`` /
	``requires_alternating_roles``) are derived from the tokenizer's own template. ``hf_id``/``dtype`` are read
	back from the model so the participant can still be pickled + re-loaded on a spawn worker. Because it
	constructs ``cls``, calling it on a subclass returns that subclass — use ``AutoModelParticipant.from_model``
	to have the family resolved from ``config.model_type``."""
	from ...loading import load_tokenizer, derive_chat_flags  # lazy: loading imports this module
	hf_id = getattr(model.config, "_name_or_path", None) or None
	if tokenizer is None:
		if not hf_id:
			raise ValueError("cannot infer a tokenizer: model.config._name_or_path is empty; pass tokenizer=")
		tokenizer = load_tokenizer(hf_id)
	supports_system, requires_alt = derive_chat_flags(tokenizer)
	p = cls(name=name, device=device, hf_id=hf_id, dtype=dtype_to_str(model.dtype),
	        attn=getattr(model, "_resolved_attn", "flash_attention_2"), **participant_kwargs)
	p._model = model
	p._tokenizer = tokenizer
	p.supports_system_role = supports_system
	p.requires_alternating_roles = requires_alt
	if p.device is None:
		p.device = model.device
	return p

@classmethod
def from_pretrained(cls, id_or_path: str | Path, *, name: str, device: str | torch.device = "cuda",
                    load_kwargs: dict | None = None, **participant_kwargs) -> Self:
	"""Build a participant of THIS class that will load ``id_or_path`` (an HF id or local path) **lazily on first
	use** — no weights are touched here. ``load_kwargs`` (``dtype`` / ``attn`` / ``quant`` / ``revision`` /
	``weights_path``) are recorded for that deferred load; ``participant_kwargs`` (``temperature``,
	``max_new_tokens``, ``system_prompt``, ``tools``, ``kv_reuse``, …) go to the participant. When the load fires
	it is process-cached, so all same-(id, device, dtype, …) participants share ONE model object. As with
	``from_model``, the class is ``cls`` — use ``AutoModelParticipant.from_pretrained`` to resolve the family
	from ``config.model_type`` (it reads only the config, still no weights)."""
	lk = dict(load_kwargs or {})
	return cls(name=name, device=device, hf_id=str(id_or_path),
	           weights_path=lk.get("weights_path"), dtype=dtype_to_str(lk.get("dtype", torch.bfloat16)),
	           attn=lk.get("attn", "flash_attention_2"), quant=lk.get("quant"), revision=lk.get("revision"),
	           **participant_kwargs)

def generate_batch(self, views: list[list[dict]], *, turn: int | None = None,
                   group_seed: int | None = None, max_new_tokens: int | None = None) -> list[Message]:
	"""Batched generation for many independent conversations that share THIS model (``throughput`` mode).

	Renders each ``view`` with this model's chat template and runs **one** ``model.generate`` over the
	left-padded batch, returning one ``Message`` per view — the co-stepping throughput win (5-20x on a
	rollout). ``max_new_tokens`` overrides this participant's per-turn cap for the batch (the co-stepper passes a
	``turn_cap`` here so a ``TokenBudget`` can shrink the final round). **No tools/steering/capture/logprobs**
	here: callers needing those fall back to the per-conversation ``generate``. Tokens are **not** guaranteed
	identical to unbatched — batch composition and the single global RNG perturb rows (see PLAN §Execution
	modes); only distributional reproducibility holds. ``metadata['batched']`` marks these turns;
	``metadata['shared_prefill']`` marks the fast path.
	"""
	if not views:
		return []
	seed = group_seed if group_seed is not None else self.seed
	if seed is not None:
		torch.manual_seed(seed)

	template_kwargs = {}
	if isinstance(self.thinking, bool):
		template_kwargs["enable_thinking"] = self.thinking
	prompts = [self.tokenizer.apply_chat_template(v, tokenize=False, add_generation_prompt=True,
	                                              **template_kwargs) for v in views]

	do_sample = bool(self.temperature and self.temperature > 0)
	gen = dict(max_new_tokens=max_new_tokens if max_new_tokens is not None else self.max_new_tokens,
	           do_sample=do_sample, pad_token_id=self.tokenizer.pad_token_id)
	if do_sample:
		gen.update(temperature=self.temperature, top_p=self.top_p)

	# Shared-scenario fast path: when every rendered prompt is token-identical (turn 1 of a rollout off one
	# scenario), prefill the shared prefix ONCE and fork to N samples via ``num_return_sequences`` instead of
	# N redundant prefills — the "prefill the shared prefix once per participant" win (PLAN §Shared-scenario
	# reuse). Only valid when sampling (identical greedy rows would be identical, so batching is pointless).
	shared_prefill = do_sample and len(views) > 1 and len(set(prompts)) == 1
	if shared_prefill:
		enc = self.tokenizer(prompts[0], return_tensors="pt", add_special_tokens=False).to(self.device)
		prompt_len = enc["input_ids"].shape[1]
		with torch.inference_mode():
			out = self.model.generate(**enc, num_return_sequences=len(views), **gen)
		new = out[:, prompt_len:]
	else:
		prev_side = self.tokenizer.padding_side
		self.tokenizer.padding_side = "left"
		try:
			enc = self.tokenizer(prompts, return_tensors="pt", padding=True,
			                     add_special_tokens=False).to(self.device)
		finally:
			self.tokenizer.padding_side = prev_side
		prompt_len = enc["input_ids"].shape[1]
		with torch.inference_mode():
			out = self.model.generate(**enc, **gen)
		new = out[:, prompt_len:]

	pad_id = self.tokenizer.pad_token_id
	messages = []
	for row in new:
		# Trailing pad ids appear once a row hits EOS before its peers; strip them for an honest token count.
		keep = row[row != pad_id] if pad_id is not None else row
		raw = self.tokenizer.decode(keep, skip_special_tokens=True)
		content, parsed_think = self.split_reasoning(raw)
		messages.append(Message(author=self.name, content=content, metadata={
			"raw_completion": raw, "parsed_think": parsed_think, "n_tokens": int(keep.shape[0]),
			"batched": True, "shared_prefill": shared_prefill,
		}))
	return messages

def parse_tool_calls(self, text: str) -> list:
	"""Parse tool calls out of a generation. Base handles the common Hermes/Qwen ``<tool_call>{json}</tool_call>``
	format; families with other formats (Gemma's ```` ```tool_code ````, Llama's ``<|python_tag|>``) override
	this. An unrecognized/absent call yields ``[]`` so the loop treats the output as a final message."""
	from ...tools.tool_call import ToolCall

	calls = []
	for match in _TOOL_CALL_RE.finditer(text):
		try:
			data = json.loads(match.group(1))
			calls.append(ToolCall(name=data["name"], arguments=data.get("arguments", {}), raw=match.group(0)))
		except (json.JSONDecodeError, KeyError) as exc:
			logger.debug("dropping malformed tool call %r: %s", match.group(0), exc)
			continue
	return calls

def render_tool_result(self, call, result) -> dict:
	"""Render a tool result as the standard structured ``tool`` message; the tokenizer's own template turns it
	into the family-native format."""
	return {"role": "tool", "name": result.name, "content": result.output}

def split_reasoning(self, text: str) -> tuple[str, str | None]:
	"""Split a raw completion into ``(visible_content, parsed_think)``. Base handles the ``<think>...</think>``
	convention; families with other delimiters override this."""
	match = _THINK_RE.match(text)
	if not match:
		return text.strip(), None
	return text[match.end():].strip(), match.group(1).strip()