nip.utils.malt_forest.visualise.MaltForestVisualiser

class nip.utils.malt_forest.visualise.MaltForestVisualiser(malt_forest: list[MaltTree], hyper_params: HyperParameters, protocol_handler: ProtocolHandler, environment_class: type[PureTextEnvironment])

Class for visualising the MALT forest of trees.

Parameters:

• malt_forest (list[MaltTree]) – The MALT forest to be visualised.

• hyper_params (HyperParameters) – The hyper-parameters of the experiment used to generate the rollouts from which the MALT forest was built.

• protocol_handler (ProtocolHandler) – The protocol handler constructed for the hyper-parameters. This will be functionally identical to the one used to generate the rollouts from which the MALT forest was built.

• environment_class (PureTextEnvironment) – The environment class determined from the hyper-parameters. If instantiated, this would be functionally identical to the one used to generate the rollouts from which the MALT forest was built.

Methods Summary

__init__(malt_forest, hyper_params, ...)
_build_dict_tree(malt_node)	Build a dictionary representation of a MALT tree, recursively.
as_jsonl()	Build a JSONL representation of the MALT forest.
as_yaml()	Build a YAML representation of the MALT forest.
build_dict_forest()	Build a dictionary representation of the MALT forest.
from_wandb_run(run_id, iteration, wandb_project)	Build a MaltForestVisualiser by from a wandb run.
visualise([format])	Visualise the MALT forest in the specified format.

Methods

__init__(malt_forest: list[MaltTree], hyper_params: HyperParameters, protocol_handler: ProtocolHandler, environment_class: type[PureTextEnvironment])

_build_dict_tree(malt_node: MaltNode) → dict

Build a dictionary representation of a MALT tree, recursively.

Parameters:

malt_node (MaltNode) – The root node of the MALT tree to be converted to a dictionary.

Returns:

malt_dict_tree (dict) – The MALT tree represented as a dictionary. Each node is represented as a dictionary with keys for the messages sent, various metadata, and the children of the node.

as_jsonl() → str

Build a JSONL representation of the MALT forest.

A JSONL (JSON Lines) file contains one JSON object per line. See JSON Lines.

Returns:

malt_jsonl_forest (str) – The MALT forest represented as a JSONL string. Each tree is represented as a nested dictionary.

as_yaml() → str

Build a YAML representation of the MALT forest.

Returns:

malt_yaml_forest (str) – The MALT forest represented as a YAML string. Each tree is represented as a nested dictionary.

build_dict_forest() → list[dict]

Build a dictionary representation of the MALT forest.

Returns:

malt_dict_forest (list[dict]) – The MALT forest represented as a list of dictionaries. Each contains the whole tree as nested dictionaries.

classmethod from_wandb_run(run_id: str, iteration: int | str, wandb_project: str, wandb_entity: str | None = None, wandb_api: Api | None = None) → Self

Build a MaltForestVisualiser by from a wandb run.

This method downloads the rollouts from the wandb run, builds the MALT forest from them, and constructs the visualiser.

Parameters:

• run_id (str) – The id of the wandb run to download the rollouts from.

• iteration (int | str) – The experiment iteration to use as the rollouts.

• wandb_project (str) – The project of the wandb run.

• wandb_entity (str, optional) – The entity of the wandb run. If not provided, the default entity will be used.

• wandb_api (WandbApi, optional) – The wandb API instance to use. If not provided, a new instance will be created.

Returns:

malt_forest_visualiser (MaltForestVisualiser) – The MALT forest visualiser constructed from the wandb run.

visualise(format: Literal['yaml', 'jsonl'] = 'yaml') → str

Visualise the MALT forest in the specified format.

Parameters:

format (Literal["yaml", "jsonl"]) – The format to visualise the MALT forest in.

Returns:

malt_forest_visualisation (str) – The MALT forest visualisation in the specified format. Each tree is represented as a nested dictionary.