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
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.