wehub-resource-sync
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
from tempfile import tempdir
|
|
from fairseq.data.encoders.gpt2_bpe import GPT2BPE, GPT2BPEConfig
|
|
from fairseq.data.encoders import register_bpe
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
INSERT_OR_REPLACE = 0 # 1 for replace and 0 for insert
|
|
|
|
@register_bpe("gpt2es", dataclass=GPT2BPEConfig) # as stands for attention space
|
|
class GPT2BPEEnhancedSpace(GPT2BPE):
|
|
def __init__(self, cfg):
|
|
logger.info('Using the GPT2BPEEnhancedSpace.')
|
|
super().__init__(cfg)
|
|
|
|
def encode(self, x: str) -> str:
|
|
# only for sroie
|
|
assert not x.startswith(' ')
|
|
assert not x.endswith(' ')
|
|
if INSERT_OR_REPLACE == 1:
|
|
temp = []
|
|
word = ''
|
|
for ch in x:
|
|
if ch == ' ':
|
|
if word:
|
|
temp.append(word)
|
|
word = ''
|
|
temp.append('<s>')
|
|
else:
|
|
word += ch
|
|
if word:
|
|
temp.append(word)
|
|
|
|
for i in range(len(temp)):
|
|
if temp[i] != '<s>':
|
|
temp[i] = ' '.join(map(str, self.bpe.encode(temp[i])))
|
|
|
|
return ' '.join(temp)
|
|
elif INSERT_OR_REPLACE == 0:
|
|
temp = []
|
|
word = ''
|
|
for ch in x:
|
|
if ch == ' ':
|
|
if word:
|
|
temp.append(word)
|
|
word = ' '
|
|
temp.append('<s>')
|
|
else:
|
|
word += ch
|
|
if word:
|
|
temp.append(word)
|
|
|
|
for i in range(len(temp)):
|
|
if temp[i] != '<s>':
|
|
temp[i] = ' '.join(map(str, self.bpe.encode(temp[i])))
|
|
|
|
return ' '.join(temp)
|
|
|
|
def decode(self, x: str) -> str:
|
|
if INSERT_OR_REPLACE == 1:
|
|
return self.bpe.decode(
|
|
[int(tok) if tok not in {"<unk>", "<mask>", "<s>"} else tok for tok in x.split()]
|
|
).replace('<s>', ' ')
|
|
elif INSERT_OR_REPLACE == 0:
|
|
return self.bpe.decode(
|
|
[int(tok) if tok not in {"<unk>", "<mask>", "<s>"} else tok for tok in x.split()]
|
|
).replace('<s>', '')
|
|
|
|
def is_beginning_of_word(self, x: str) -> bool:
|
|
return self.decode(x).startswith(" ")
|
|
|