File size: 3,535 Bytes
409d4fb | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """Text frontend: NFKD normalisation, light cleanup, "<lang>text</lang>" tagging, and a
code-point -> id lookup (unicode_indexer.json, 8322 ids)."""
import json
import re
from unicodedata import normalize
AVAILABLE_LANGS = ["en", "ko", "ja", "ar", "bg", "cs", "da", "de", "el", "es", "et", "fi", "fr", "hi", "hr", "hu", "id", "it", "lt", "lv", "nl", "pl", "pt", "ro", "ru", "sk", "sl", "sv", "tr", "uk", "vi", "na"]
_EMOJI = re.compile(
"[\U0001f600-\U0001f64f\U0001f300-\U0001f5ff\U0001f680-\U0001f6ff\U0001f700-\U0001f77f"
"\U0001f780-\U0001f7ff\U0001f800-\U0001f8ff\U0001f900-\U0001f9ff\U0001fa00-\U0001fa6f"
"\U0001fa70-\U0001faff☀-⛿✀-➿\U0001f1e6-\U0001f1ff]+",
flags=re.UNICODE,
)
_REPL = {"–": "-", "‑": "-", "—": "-", "_": " ", "“": '"', "”": '"', "‘": "'", "’": "'",
"´": "'", "`": "'", "[": " ", "]": " ", "|": " ", "/": " ", "#": " ", "→": " ", "←": " "}
_EXPR = {"@": " at ", "e.g.,": "for example, ", "i.e.,": "that is, "}
def preprocess_text(text: str, lang: str) -> str:
text = normalize("NFKD", text)
text = _EMOJI.sub("", text)
for k, v in _REPL.items():
text = text.replace(k, v)
text = re.sub(r"[♥☆♡©\\]", "", text)
for k, v in _EXPR.items():
text = text.replace(k, v)
for p in [",", r"\.", "!", r"\?", ";", ":", "'"]:
text = re.sub(r" " + p, p.replace("\\", ""), text)
while '""' in text:
text = text.replace('""', '"')
while "''" in text:
text = text.replace("''", "'")
while "``" in text:
text = text.replace("``", "`")
text = re.sub(r"\s+", " ", text).strip()
if not re.search(r"[.!?;:,'\"')\]}…。」』】〉》›»]$", text):
text += "."
if lang not in AVAILABLE_LANGS:
raise ValueError(f"Invalid language: {lang}")
return f"<{lang}>" + text + f"</{lang}>"
class TextProcessor:
def __init__(self, unicode_indexer_path: str):
self.indexer_path = unicode_indexer_path
with open(unicode_indexer_path) as f:
self.indexer = json.load(f) # list: code point -> id (-1 = unknown)
def encode(self, text: str, lang: str, preprocess: bool = True) -> list[int]:
s = preprocess_text(text, lang) if preprocess else text
return [self.indexer[ord(c)] if ord(c) < len(self.indexer) else -1 for c in s]
def unknown_chars(self, text: str, lang: str) -> set[str]:
s = preprocess_text(text, lang)
return {c for c in s if ord(c) >= len(self.indexer) or self.indexer[ord(c)] < 0}
def batch_numpy(self, texts, langs):
"""ids (B, T) int64 and mask (B, 1, T) float32 as numpy arrays."""
import numpy as np
seqs = [self.encode(t, l) for t, l in zip(texts, langs)]
T = max(len(s) for s in seqs)
ids = np.zeros((len(seqs), T), np.int64)
mask = np.zeros((len(seqs), 1, T), np.float32)
for i, s in enumerate(seqs):
ids[i, : len(s)] = s
mask[i, :, : len(s)] = 1
return ids, mask
def batch(self, texts: list[str], langs: list[str], device=None):
import torch
seqs = [self.encode(t, l) for t, l in zip(texts, langs)]
lengths = torch.tensor([len(s) for s in seqs])
ids = torch.zeros(len(seqs), int(lengths.max()), dtype=torch.long)
for i, s in enumerate(seqs):
ids[i, : len(s)] = torch.tensor(s)
mask = (torch.arange(ids.shape[1])[None] < lengths[:, None]).float().unsqueeze(1)
return ids.to(device), mask.to(device)
|