| """ |
| Dataset and Tokenization module for Part 1 (Vietnamese/Japanese to English Translation). |
| Handles belumind/en-vi-ja-curated-500k-triplets dataset. |
| Supports Byte-level BPE tokenization handling non-romanized Japanese and Vietnamese characters. |
| """ |
|
|
| import os |
| from typing import Dict, List, Tuple, Optional, Any |
| import torch |
| from torch.utils.data import Dataset, DataLoader |
| from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors |
| from transformers import PreTrainedTokenizerFast |
|
|
|
|
| SPECIAL_TOKENS = ["<pad>", "<unk>", "<bos>", "<eos>", "<vi>", "<ja>", "<en>"] |
|
|
|
|
| def build_or_load_tokenizer( |
| vocab_size: int = 16000, |
| save_path: str = "assets/tokenizer_multilingual.json", |
| corpus_texts: Optional[List[str]] = None |
| ) -> PreTrainedTokenizerFast: |
| """ |
| Builds or loads a Byte-level BPE tokenizer capable of processing |
| arbitrary UTF-8 bytes (Vietnamese diacritics and Japanese Kanji/Kana). |
| """ |
| if os.path.exists(save_path): |
| tokenizer = Tokenizer.from_file(save_path) |
| fast_tok = PreTrainedTokenizerFast( |
| tokenizer_object=tokenizer, |
| bos_token="<bos>", |
| eos_token="<eos>", |
| pad_token="<pad>", |
| unk_token="<unk>", |
| ) |
| return fast_tok |
|
|
| |
| raw_tokenizer = Tokenizer(models.BPE(unk_token="<unk>")) |
| raw_tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| raw_tokenizer.decoder = decoders.ByteLevel() |
|
|
| trainer = trainers.BpeTrainer( |
| vocab_size=vocab_size, |
| special_tokens=SPECIAL_TOKENS, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), |
| show_progress=False |
| ) |
|
|
| if corpus_texts is not None and len(corpus_texts) > 0: |
| raw_tokenizer.train_from_iterator(corpus_texts, trainer=trainer) |
| else: |
| |
| seed_data = [ |
| "Hello world, this is English.", |
| "こんにちは世界、日本語のテストです。", |
| "Xin chào thế giới, đây là kiểm tra tiếng Việt.", |
| ] * 100 |
| raw_tokenizer.train_from_iterator(seed_data, trainer=trainer) |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True) |
| raw_tokenizer.save(save_path) |
|
|
| fast_tok = PreTrainedTokenizerFast( |
| tokenizer_object=raw_tokenizer, |
| bos_token="<bos>", |
| eos_token="<eos>", |
| pad_token="<pad>", |
| unk_token="<unk>", |
| ) |
| return fast_tok |
|
|
|
|
| class TranslationDataset(Dataset): |
| """ |
| Dataset creating Vietnamese->English and Japanese->English samples. |
| Each triplet (vi, ja, en) produces two training sequences: |
| 1. <vi> {vi_text} <en> {en_text} <eos> |
| 2. <ja> {ja_text} <en> {en_text} <eos> |
| """ |
| def __init__( |
| self, |
| samples: List[Dict[str, str]], |
| tokenizer: PreTrainedTokenizerFast, |
| max_seq_len: int = 256, |
| loss_on_target_only: bool = True |
| ): |
| self.tokenizer = tokenizer |
| self.max_seq_len = max_seq_len |
| self.loss_on_target_only = loss_on_target_only |
| self.data: List[Dict[str, Any]] = [] |
|
|
| for item in samples: |
| en = item.get("en", "").strip() |
| vi = item.get("vi", "").strip() |
| ja = item.get("ja", "").strip() |
|
|
| if vi and en: |
| self.data.append({ |
| "src_lang": "vi", |
| "src_text": vi, |
| "tgt_text": en, |
| "prompt": f"<vi> {vi} <en>", |
| "full_text": f"<vi> {vi} <en> {en} <eos>" |
| }) |
| if ja and en: |
| self.data.append({ |
| "src_lang": "ja", |
| "src_text": ja, |
| "tgt_text": en, |
| "prompt": f"<ja> {ja} <en>", |
| "full_text": f"<ja> {ja} <en> {en} <eos>" |
| }) |
|
|
| def __len__(self) -> int: |
| return len(self.data) |
|
|
| def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: |
| item = self.data[idx] |
| full_text = item["full_text"] |
| prompt = item["prompt"] |
|
|
| full_encoding = self.tokenizer( |
| full_text, |
| max_length=self.max_seq_len, |
| padding="max_length", |
| truncation=True, |
| return_tensors="pt" |
| ) |
| input_ids = full_encoding["input_ids"].squeeze(0) |
|
|
| targets = input_ids.clone() |
| pad_token_id = self.tokenizer.pad_token_id or 0 |
| targets[targets == pad_token_id] = -100 |
|
|
| if self.loss_on_target_only: |
| prompt_encoding = self.tokenizer( |
| prompt, |
| max_length=self.max_seq_len, |
| truncation=True, |
| return_tensors="pt" |
| ) |
| prompt_len = min(prompt_encoding["input_ids"].shape[1], self.max_seq_len) |
| targets[:prompt_len] = -100 |
|
|
| return { |
| "input_ids": input_ids, |
| "targets": targets, |
| "src_lang": item["src_lang"], |
| "src_text": item["src_text"], |
| "tgt_text": item["tgt_text"], |
| "prompt": item["prompt"] |
| } |
|
|
|
|
| def collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]: |
| input_ids = torch.stack([b["input_ids"] for b in batch]) |
| targets = torch.stack([b["targets"] for b in batch]) |
| src_langs = [b["src_lang"] for b in batch] |
| src_texts = [b["src_text"] for b in batch] |
| tgt_texts = [b["tgt_text"] for b in batch] |
| prompts = [b["prompt"] for b in batch] |
|
|
| return { |
| "input_ids": input_ids, |
| "targets": targets, |
| "src_langs": src_langs, |
| "src_texts": src_texts, |
| "tgt_texts": tgt_texts, |
| "prompts": prompts |
| } |
|
|
