Spaces:
Running
Running
File size: 1,320 Bytes
03f6091 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# -*- coding: utf-8 -*-
import torch
from transformers import AutoTokenizer
from torchnlp.encoders.text.text_encoder import TextEncoder
from .tokenizer_base import TextEncoderBase
class HFTextEncoder(TextEncoderBase):
"""
Wrapper arround transformers AutoTokenizer:
- https://huggingface.co/transformers/model_doc/auto.html#autotokenizer
:param model: model to be used.
"""
def __init__(self, model: str) -> None:
super().__init__()
self.tokenizer = AutoTokenizer.from_pretrained(model)
# Properties from the base class
self.stoi = self.tokenizer.get_vocab()
self.itos = {v: k for k, v in self.stoi.items()}
self._bos_index = self.tokenizer.cls_token_id
self._pad_index = self.tokenizer.pad_token_id
self._eos_index = self.tokenizer.sep_token_id
self._unk_index = self.tokenizer.unk_token_id
self._mask_index = self.tokenizer.mask_token_id
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 torch.tensor(self.tokenizer(sequence, truncation=False)["input_ids"])
|