Spaces:
Sleeping
Sleeping
File size: 1,587 Bytes
c1596ac | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import json
from collections import Counter
import re
import sentencepiece as spm
def tokenizer(captions):
text = captions.lower()
text = re.sub(r"([.,!?])", r" \1 ", text) # 특수문자 제거
tokens = text.split()
return tokens
def sub_tokenizer(caption, sp):
tokens = sp.encode(caption, out_type=str)
return tokens
def build_vocab(json_path, min_freq=3, max_size=10000, use_subword=False, sp_model_path="/workspace/src/dataset/sub_tokenizer.model"):
w2i = dict()
i2w = dict()
# ==================================================
# SentencePiece tokenizer 사용
# ==================================================
if use_subword:
sp = spm.SentencePieceProcessor()
sp.load(sp_model_path)
voca_size = sp.get_piece_size()
for i in range(voca_size):
token = sp.id_to_piece(i)
w2i[token] = i
i2w[i] = token
else:
with open(json_path, 'r') as f:
data = json.load(f)
counter = Counter()
for item in data:
captions = item["captions"]
for caption in captions:
tokens = tokenizer(caption)
counter.update(tokens)
words = [w for w, freq in counter.most_common() if freq >= min_freq]
voca = ["<pad>", "<sos>", "<eos>", "<unk>"]
voca.extend(words[:max_size-4])
voca_size = len(voca)
for i, w in enumerate(voca):
w2i[w] = i
i2w[i] = w
print(voca_size)
return w2i, i2w, voca_size |