interlens.participant.participants.model_participant¶

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()