Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| from typing import Callable, Dict | |
| import torch | |
| from torchnlp.encoders.text.text_encoder import TextEncoder | |
| from .tokenizer_base import TextEncoderBase | |
| class XLMRTextEncoder(TextEncoderBase): | |
| """ | |
| XLM-RoBERTa encoder from Fairseq. | |
| :param tokenizer_func: XLM tokenization function. | |
| This can be easily obtain from the fairseq model (e.g: xlmr.encode callable) | |
| :param vocabulary: the dictionary containing the XLM-R vocabulary. | |
| This can be easily obtain from the fairseq model | |
| (e.g: xlmr.task.source_dictionary.__dict__['indices']) | |
| """ | |
| def __init__(self, encode_func: Callable, vocabulary: Dict[str, int]) -> None: | |
| super().__init__() | |
| self.encode_func = encode_func | |
| # Properties from the base class | |
| self.stoi = vocabulary | |
| self.itos = {v: k for k, v in vocabulary.items()} | |
| self._pad_index = self.stoi["<pad>"] | |
| self._eos_index = self.stoi["</s>"] | |
| self._unk_index = self.stoi["<unk>"] | |
| self._bos_index = self.stoi["<s>"] | |
| self._mask_index = self.stoi["<mask>"] | |
| def encode(self, sequence: str) -> torch.Tensor: | |
| """Encodes a 'sequence'. | |
| :param sequence: String 'sequence' to encode. | |
| Returns: | |
| - torch.Tensor: Encoding of the 'sequence'. | |
| """ | |
| sequence = TextEncoder.encode(self, sequence) | |
| return self.encode_func(sequence) | |