CWE

class minicons.cwe.CWE(model_name: str, device: str | None = 'cpu', pretrained: bool = True, **kwargs)

Bases: object

Implements the contextualized word embedding class to facilitate extraction of word representations from a given transformer model. [TEST]

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: str | List[str], layer: 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: List[List[str | Tuple[int, int]]] | List[str | Tuple[int, int]], layer: int | List[int] = None, multi_strategy='last') 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: str | List[str], layer: int | List[int], pooler: str = 'mean') 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: Tuple[str] | List[Tuple[str]], layer: int = None, multi_strategy='last') Tuple

Extract representations of pairs of words from a given sentence from the model at a given layer.

Parameters:
  • text (Union[Tuple[str], List[Tuple[str]]]) – Input consisting of [(sentence, word1, word2)], where sentence is an input sentence, and word1, word2 are two words present in the sentence that will be masked out.

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

Returns:

Tuple consisting of torch tensors or lists of torch tensors corresponding to word representations