|
|
import os |
|
|
import re |
|
|
import pickle |
|
|
from typing import List, Optional, Tuple |
|
|
from transformers import PreTrainedTokenizer |
|
|
from tokenizers import Tokenizer |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
try: |
|
|
import ahocorasick |
|
|
except ImportError: |
|
|
raise ImportError( |
|
|
"This tokenizer requires the 'pyahocorasick' library. " |
|
|
"Please install it using: pip install pyahocorasick" |
|
|
) |
|
|
|
|
|
class MiLLiTokenizer(PreTrainedTokenizer): |
|
|
"""miLLi 1.0: Azerbaijani Hybrid Tokenizer with Phonological Restoration""" |
|
|
|
|
|
model_input_names = ["input_ids", "attention_mask"] |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
repo_id="elshadrahimov/miLLi-1.0", |
|
|
trie_file="milli_trie_v1.pkl", |
|
|
bpe_file="milli_bpe_v1.json", |
|
|
unk_token="<UNK>", |
|
|
pad_token="<PAD>", |
|
|
eos_token="<EOS>", |
|
|
**kwargs |
|
|
): |
|
|
local_dir = os.path.dirname(__file__) |
|
|
local_trie_path = os.path.join(local_dir, trie_file) |
|
|
local_bpe_path = os.path.join(local_dir, bpe_file) |
|
|
|
|
|
if os.path.exists(local_trie_path): |
|
|
self.trie_path = local_trie_path |
|
|
else: |
|
|
try: |
|
|
self.trie_path = hf_hub_download(repo_id=repo_id, filename=trie_file) |
|
|
except Exception as e: |
|
|
raise FileNotFoundError(f"Could not download Trie file from Hub: {e}") |
|
|
|
|
|
if os.path.exists(local_bpe_path): |
|
|
self.bpe_path = local_bpe_path |
|
|
else: |
|
|
try: |
|
|
self.bpe_path = hf_hub_download(repo_id=repo_id, filename=bpe_file) |
|
|
except Exception as e: |
|
|
raise FileNotFoundError(f"Could not download BPE file from Hub: {e}") |
|
|
|
|
|
with open(self.trie_path, "rb") as f: |
|
|
self.trie = pickle.load(f) |
|
|
|
|
|
self.bpe_backend = Tokenizer.from_file(self.bpe_path) |
|
|
|
|
|
super().__init__( |
|
|
unk_token=unk_token, |
|
|
pad_token=pad_token, |
|
|
eos_token=eos_token, |
|
|
**kwargs |
|
|
) |
|
|
|
|
|
@property |
|
|
def vocab_size(self) -> int: |
|
|
return self.bpe_backend.get_vocab_size() |
|
|
|
|
|
def get_vocab(self): |
|
|
return self.bpe_backend.get_vocab() |
|
|
|
|
|
def _normalize_phonology(self, sub_word: str) -> Optional[str]: |
|
|
if "_" + sub_word in self.trie: |
|
|
return sub_word |
|
|
|
|
|
replacements = {'ğ': 'q', 'y': 'k'} |
|
|
if sub_word and sub_word[-1] in replacements: |
|
|
restored = sub_word[:-1] + replacements[sub_word[-1]] |
|
|
if "_" + restored in self.trie: |
|
|
return restored |
|
|
|
|
|
vowels = ['ı', 'i', 'u', 'ü'] |
|
|
if len(sub_word) >= 2: |
|
|
for v in vowels: |
|
|
restored = sub_word[:-1] + v + sub_word[-1] |
|
|
if "_" + restored in self.trie: |
|
|
return restored |
|
|
return None |
|
|
|
|
|
def _milli_pre_tokenize(self, text: str) -> List[str]: |
|
|
raw_tokens = re.findall(r'\w+|[^\w\s]', text, re.UNICODE) |
|
|
processed_parts = [] |
|
|
for token in raw_tokens: |
|
|
if not re.match(r'\w+', token): |
|
|
processed_parts.append(token); continue |
|
|
|
|
|
is_upper = token[0].isupper() |
|
|
word = token.lower() |
|
|
if is_upper: processed_parts.append("<UPPER>") |
|
|
|
|
|
found_root, best_root_len = None, 0 |
|
|
for i in range(len(word), 0, -1): |
|
|
root = self._normalize_phonology(word[:i]) |
|
|
if root: |
|
|
found_root = root; best_root_len = i; break |
|
|
|
|
|
if found_root: |
|
|
processed_parts.append("_" + found_root) |
|
|
remainder = word[best_root_len:] |
|
|
if remainder: processed_parts.append("##" + remainder) |
|
|
else: |
|
|
processed_parts.append("_" + word) |
|
|
return processed_parts |
|
|
|
|
|
def _tokenize(self, text: str) -> List[str]: |
|
|
pre_parts = self._milli_pre_tokenize(text) |
|
|
final_tokens = [] |
|
|
for part in pre_parts: |
|
|
clean_part = part.replace('_', '').replace('##', '') |
|
|
if not re.match(r'\w+', clean_part): |
|
|
final_tokens.append(part); continue |
|
|
encoding = self.bpe_backend.encode(part) |
|
|
final_tokens.extend(encoding.tokens) |
|
|
return final_tokens |
|
|
|
|
|
def _convert_token_to_id(self, token: str) -> int: |
|
|
return self.bpe_backend.token_to_id(token) or self.bpe_backend.token_to_id(self.unk_token) |
|
|
|
|
|
def _convert_id_to_token(self, index: int) -> str: |
|
|
return self.bpe_backend.id_to_token(index) |
|
|
|
|
|
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: |
|
|
if not os.path.isdir(save_directory): |
|
|
os.makedirs(save_directory) |
|
|
vocab_file = os.path.join(save_directory, (filename_prefix or "") + "milli_bpe_v1.json") |
|
|
self.bpe_backend.save(vocab_file) |
|
|
return (vocab_file,) |
