minicons.cwe module

class minicons.cwe.CWE(model_name: str, device: Optional[str] = 'cpu', pretrained: bool = True)

Bases: object

Implements the contextualized word embedding class to facilitate extraction of word representations form a given transformer model.

Parameters

  • model_name (str) – name of the model, should either be a path to a model (.pt or .bin file) stored locally, or a pretrained model stored on the Huggingface Model Hub.

  • device (str, optional) – device type that the model should be loaded on, options: cpu or cuda:{0, 1, …}

  • pretrained (bool) – whether to load the model with pretrained weights. loads a randomly initialized model if False. Default = True.

encode_text(text: Union[str, List[str]], layer: Optional[Union[int, List[int]]] = None) → Tuple

Encodes batch of raw sentences using the model to return hidden states at a given layer.

Parameters

  • text (Union[str, List[str]]) – batch of raw sentences

  • layer (int) – layer from which the hidden states are extracted.

Returns

Tuple (input_ids, hidden_states)

extract_representation(sentence_words: Union[List[List[Union[str, Tuple[int, int]]]], List[Union[str, Tuple[int, int]]]], layer: Optional[Union[int, List[int]]] = None) → Union[torch.Tensor, List[torch.Tensor]]

Extract representations from the model at a given layer.

Parameters

  • sentence_words (Union[List[List[Union[str, Union[Tuple(int, int), str]]]], List[Union[str, Union[Tuple(int, int), str]]]]) – Input consisting of [(sentence, word)], where sentence is an input sentence, and word is a word present in the sentence that will be masked out, or [(sentence, (start, end))], where (start, end) is a tuple consisting of the character span indices that form the word.

  • layer (Union[int, List[int]]) – layer(s) from which the hidden states are extracted.

Returns

torch tensors or list of torch tensors corresponding to word representations

extract_sentence_representation(sentences: Union[str, List[str]], layer: Union[int, List[int]], pooler: str = 'mean') → Union[torch.Tensor, List[torch.Tensor]]

Extract representations of input sentences from one or more layers in the model.

Parameters

  • sentences (Union[str, List[str]]) – Input consisting of one or more sentences.

  • layer (Union[int, List[int]]) – one or more layer from which the representations are extracted.

  • pooler (str) – pooling logic (mean, max, min)

Returns

torch.Tensor

extract_paired_representations(sentence_words: Union[Tuple[str], List[Tuple[str]]], layer: Optional[int] = None) → Tuple