extract

Extractor

class diagnnose.extract.extractor.Extractor(model: LanguageModel, corpus: ~diagnnose.corpus.corpus.Corpus, activation_names: ~typing.Optional[~typing.List[~typing.Tuple[int, str]]] = None, activations_dir: ~typing.Optional[str] = None, selection_func: ~typing.Union[~typing.Callable[[int, ~torchtext.data.example.Example], bool], str] = <function return_all>, batch_size: int = 1024)[source]

Bases: object

Extracts all intermediate activations of a LM from a corpus.

Only activations that are provided in activation_names will be stored in a pickle file. Each activation is written to its own file.

Parameters
  • model (LanguageModel) – Language model that inherits from LanguageModel.

  • corpus (Corpus) – Corpus containing sentences to be extracted.

  • activation_names (List[tuple[int, str]], optional) – List of (layer, activation_name) tuples. If not provided all activation_names corresponding to the model will be extracted.

  • activations_dir (str, optional) – Directory to which activations will be written. If not provided the extract() method will only return the activations without writing them to disk.

  • selection_func (Union[SelectionFunc, str]) – Function which determines if activations for a token should be extracted or not. Can also be provided as a string, indicating the method name of one of the default selection_funcs in diagnnose.activations.selection_funcs.

  • batch_size (int, optional) – Amount of sentences processed per forward step. Higher batch size increases extraction speed, but should be done accordingly to the amount of available RAM. Defaults to 1.

extract() ActivationReader[source]

Extracts embeddings from a corpus.

Uses contextlib.ExitStack to write to multiple files simultaneously.

Returns

activation_reader – After extraction an activation_reader is returned that provides direct access to the extracted activations.

Return type

ActivationReader

set_activation_ranges() None[source]

Simple Extract

diagnnose.extract.simple_extract.simple_extract(model: LanguageModel, corpus: ~diagnnose.corpus.corpus.Corpus, activation_names: ~typing.List[~typing.Tuple[int, str]], activations_dir: ~typing.Optional[str] = None, batch_size: int = 1024, selection_func: ~typing.Callable[[int, ~torchtext.data.example.Example], bool] = <function return_all>) Tuple[ActivationReader, Callable[[], None]][source]

Basic extraction method.

Parameters
  • model (LanguageModel) – Language model that inherits from LanguageModel.

  • corpus (Corpus) – Corpus containing sentences to be extracted.

  • activation_names (List[tuple[int, str]]) – List of (layer, activation_name) tuples

  • activations_dir (str, optional) – Directory to which activations will be written. If not provided the extract() method will only return the activations without writing them to disk.

  • selection_func (SelectionFunc) – Function which determines if activations for a token should be extracted or not.

  • batch_size (int, optional) – Amount of sentences processed per forward step. Higher batch size increases extraction speed, but should be done accordingly to the amount of available RAM. Defaults to 1.

Returns

  • activation_reader (ActivationReader) – ActivationReader for the activations that have been extracted.

  • remove_activations (RemoveCallback) – Callback function that can be executed at the end of a procedure that depends on the extracted activations. Removes all the activations that have been extracted. Takes no arguments.