nnsight.modeling

class nnsight.modeling.language.LanguageModel(*args, config: ~transformers.configuration_utils.PretrainedConfig | None = None, tokenizer: ~transformers.tokenization_utils.PreTrainedTokenizer | None = None, automodel: ~typing.Type[~transformers.models.auto.modeling_auto.AutoModel] = <class 'transformers.models.auto.modeling_auto.AutoModelForCausalLM'>, import_edits: bool | str = False, **kwargs)

LanguageModels are NNsight wrappers around transformers language models.

Inputs can be in the form of:

Prompt: (str) Prompts: (List[str]) Batched prompts: (List[List[str]]) Tokenized prompt: (Union[List[int], torch.Tensor]) Tokenized prompts: (Union[List[List[int]], torch.Tensor]) Direct input: (Dict[str,Any])

If using a custom model, you also need to provide the tokenizer like LanguageModel(custom_model, tokenizer=tokenizer)

Calls to generate pass arguments downstream to GenerationMixin.generate()

config

Huggingface config file loaded from repository or checkpoint.

Type:

PretrainedConfig

tokenizer

Tokenizer for LMs.

Type:

PreTrainedTokenizer

automodel

AutoModel type from transformer auto models.

Type:

Type

model

Meta version of underlying auto model.

Type:

PreTrainedModel

class Generator

class Streamer(*args, **kwargs)

export_edits(name: str | None = None, export_dir: str | None = None, variant: str = '__default__')

TODO

Parameters:

• name (Optional[str], optional) – _description_. Defaults to None.

• export_dir (Optional[str], optional) – _description_. Defaults to None.

• variant (str, optional) – _description_. Defaults to ‘__default__’.

import_edits(name: str | None = None, export_dir: str | None = None, variant: str = '__default__')

TODO

Parameters:

• name (Optional[str], optional) – _description_. Defaults to None.

• export_dir (Optional[str], optional) – _description_. Defaults to None.

• variant (str, optional) – _description_. Defaults to ‘__default__’.

class nnsight.modeling.mixins.loadable.LoadableMixin(*args, rename: Dict[str, str] | None = None, **kwargs)

class nnsight.modeling.mixins.meta.MetaMixin(*args, dispatch: bool = False, meta_buffers: bool = True, rename: Dict[str, str] | None = None, **kwargs)

class nnsight.modeling.mixins.remoteable.RemoteInterleavingTracer(fn: Callable, model: Any, *args, backend: Backend = None, **kwargs)

class nnsight.modeling.mixins.remoteable.RemoteTracer(*args, backend: Backend = None, _info: Info = None, **kwargs)

class nnsight.modeling.mixins.remoteable.RemoteableMixin(*args, dispatch: bool = False, meta_buffers: bool = True, rename: Dict[str, str] | None = None, **kwargs)

trace(*inputs: Any, backend: Backend | str | None = None, remote: bool = False, blocking: bool = True, **kwargs: Dict[str, Any])

Create a tracer for this module.

This method returns a tracer that can be used to capture and modify the execution of the module.

Example

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> with model.trace("Hello World"):
>>>     model.transformer.h[0].attn.output[0][:] = 0

>>>     output = model.output.save()
>>> print(output)

Parameters:

• *args – Arguments to pass to the tracer

• **kwargs – Keyword arguments to pass to the tracer

Returns:

An InterleavingTracer for this module

class nnsight.modeling.mixins.remoteable.StreamTracer(frame: FrameType, *args, **kwargs)

execute(fn: Callable)

Execute the compiled function.

Runs the compiled function with the necessary context to execute the traced code block.

Parameters:

fn – The compiled function to execute