Upload 6 files

src/config.py

@@ -0,0 +1,34 @@
+class Config:
+    def __init__(self):
+        self.encoder_layer = 6
+        self.decoder_layer = 6
+        self.embed_dim = 512
+        self.num_heads = 8
+        self.drop_out = 0.1
+        self.max_len = 256
+        self.vocab_size = 30_000
+        self.wmt_zh_en_path = "data/wmt_zh_en_training_corpus.csv"
+        self.tokenizer_file = "checkpoints/tokenizer.json"
+        self.batch_size = 64
+        self.compile = False
+        self.seed = 42
+        self.val_ratio = 0.1
+        self.num_workers = 4
+        self.pin_memory = True
+        self.tensorboard_dir = "log/tensorboard"
+        self.checkpoint_dir = "log/checkpoint"
+
+        self.base_lr = 3e-4
+        self.betas = (0.9, 0.98)
+        self.eps = 1e-9
+        self.weight_decay = 0.1
+        self.warmup_ratio = 0.005
+        self.start_factor = 1e-3
+        self.end_factor = 1.0
+        self.eta_min = 3e-6
+        self.max_epochs = 1
+
+        self.data_cache_dir = "data/cache.pickle"
+        self.use_cache = True
+
+        self.every_n_train_steps = 10000

src/dataset.py

@@ -0,0 +1,60 @@
+from torch.utils.data import Dataset
+import csv
+from tokenizers import Tokenizer
+import torch
+import os
+import pickle
+from src.config import Config
+
+
+class TranslateDataset(Dataset):
+    def __init__(self, config: Config):
+        super().__init__()
+        self.config = config
+        self.tokenizer: Tokenizer = Tokenizer.from_file(config.tokenizer_file)
+        self.pad_id = self.tokenizer.token_to_id("[PAD]")
+        self.pairs = []
+        if os.path.exists(config.data_cache_dir) and config.use_cache:
+            with open(config.data_cache_dir, "rb") as f:
+                self.pairs = pickle.load(f)
+        else:
+            with open(self.config.wmt_zh_en_path, mode="r", encoding="utf-8") as f:
+                reader = csv.DictReader(f)
+                for line in reader:
+                    self.pairs.append((line["0"], line["1"]))
+            if config.use_cache:
+                with open(config.data_cache_dir, "wb") as cache_f:
+                    pickle.dump(self.pairs, cache_f)
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def encode(self, text):
+        ids = self.tokenizer.encode(text).ids
+
+        if len(ids) > self.config.max_len:
+            ids = ids[: self.config.max_len]
+
+        pad_len = self.config.max_len - len(ids)
+
+        if pad_len > 0:
+            ids = ids + [self.pad_id] * pad_len
+        pad_mask = [False if i == self.pad_id else True for i in ids]
+        return torch.tensor(ids, dtype=torch.long), torch.tensor(
+            pad_mask, dtype=torch.bool
+        )
+
+    def __getitem__(self, idx):
+        zh, en = self.pairs[idx]
+
+        zh_id, zh_pad = self.encode(zh)
+
+        en_id, en_pad = self.encode(en)
+
+        return dict(
+            src=zh_id,
+            src_pad_mask=zh_pad,
+            tgt=en_id[:-1],
+            tgt_pad_mask=en_pad[:-1],
+            label=en_id[1:],
+        )

src/model.py

@@ -0,0 +1,216 @@
+from torch import nn as nn
+import torch
+from src.config import Config
+
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, drop_out=0.1):
+        super().__init__()
+
+        assert embed_dim % num_heads == 0
+
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout = nn.Dropout(drop_out)
+        self.scale = (embed_dim // num_heads) ** 0.5
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mask: torch.Tensor = None,
+        pad_mask: torch.Tensor = None,
+    ):
+        bs = q.shape[0]
+        q_len = q.shape[1]
+        k_len = k.shape[1]
+
+        Q: torch.Tensor = self.q_proj(q)
+        K: torch.Tensor = self.k_proj(k)
+        V: torch.Tensor = self.v_proj(v)
+
+        q_state = Q.view(bs, q_len, self.num_heads, -1).transpose(1, 2)
+        k_state = K.view(bs, k_len, self.num_heads, -1).transpose(1, 2)
+        v_state = V.view(bs, k_len, self.num_heads, -1).transpose(1, 2)
+
+        attn = q_state @ k_state.transpose(
+            -1, -2
+        )  # [bs, head, q_len, dim] @ [bs, head, dim, k_len] = [bs, head, q_len, k_len]
+        attn: torch.Tensor = attn / self.scale
+
+        if mask is not None:
+            attn = attn.masked_fill(~mask, -1e8)
+
+        if pad_mask is not None:
+            attn = attn.masked_fill(~pad_mask.unsqueeze(1).unsqueeze(2), -1e8)
+
+        attn = torch.softmax(attn, dim=-1)
+
+        attn = self.dropout(attn)
+
+        out = (
+            attn @ v_state
+        )  # [bs, head, q_len, k_len] @ [bs, head, k_len, dim] = [bs, head, q_len, dim]
+
+        out = out.transpose(1, 2).contiguous().view(bs, q_len, -1)
+
+        out = self.out_proj(out)
+
+        return out
+
+
+class FFN(nn.Module):
+    def __init__(self, embed_dim, drop_out=0.1):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * 4),
+            nn.ReLU(),
+            nn.Linear(embed_dim * 4, embed_dim),
+            nn.Dropout(drop_out),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, embed_dim, num_heads, drop_out=0.1):
+        super().__init__()
+        self.mha = MultiHeadAttention(
+            embed_dim=embed_dim, num_heads=num_heads, drop_out=drop_out
+        )
+        self.ffn = FFN(embed_dim=embed_dim, drop_out=drop_out)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor, pad_mask=None):
+        x = x + self.mha(x, x, x, pad_mask=pad_mask)
+        x = self.norm1(x)
+
+        x = x + self.ffn(x)
+
+        x = self.norm2(x)
+
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, config: Config):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(config.encoder_layer):
+            self.layers.append(
+                EncoderLayer(config.embed_dim, config.num_heads, config.drop_out)
+            )
+
+    def forward(self, x, pad_mask=None):
+        for layer in self.layers:
+            x = layer(x, pad_mask)
+        return x
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, embed_dim, num_heads, drop_out=0.1):
+        super().__init__()
+        self.self_attn = MultiHeadAttention(
+            embed_dim=embed_dim, num_heads=num_heads, drop_out=drop_out
+        )
+        self.cross_attn = MultiHeadAttention(
+            embed_dim=embed_dim, num_heads=num_heads, drop_out=drop_out
+        )
+        self.ffn = FFN(embed_dim=embed_dim, drop_out=drop_out)
+        self.norm0 = nn.LayerNorm(embed_dim)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor, memory, src_pad_mask=None, tgt_pad_mask=None):
+
+        x_len = x.shape[1]
+        mask = torch.ones(size=(1, 1, x_len, x_len), device=x.device, dtype=torch.bool).tril()
+
+        x = x + self.self_attn(x, x, x, mask=mask, pad_mask=tgt_pad_mask)
+
+        x = self.norm0(x)
+
+        x = x + self.cross_attn(x, memory, memory, pad_mask=src_pad_mask)
+        x = self.norm1(x)
+
+        x = x + self.ffn(x)
+
+        x = self.norm2(x)
+
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(self, config: Config):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(config.decoder_layer):
+            self.layers.append(
+                DecoderLayer(config.embed_dim, config.num_heads, config.drop_out)
+            )
+
+    def forward(self, x: torch.Tensor, memory, src_pad_mask=None, tgt_pad_mask=None):
+        for layer in self.layers:
+            x = layer(x, memory, src_pad_mask=src_pad_mask, tgt_pad_mask=tgt_pad_mask)
+        return x
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(self, config: Config):
+        super().__init__()
+        pe = torch.zeros(config.max_len, config.embed_dim)
+        pos = torch.arange(0, config.max_len, 1).float().unsqueeze(1)
+        _2i = torch.arange(0, config.embed_dim, 2)
+        pe[:, 0::2] = torch.sin(pos / (10000 ** (_2i / config.embed_dim)))
+        pe[:, 1::2] = torch.cos(pos / (10000 ** (_2i / config.embed_dim)))
+
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x_len = x.shape[1]
+        return x + self.pe[:, :x_len].to(dtype=x.dtype)
+
+
+class TranslateModel(nn.Module):
+    def __init__(self, config: Config):
+        super().__init__()
+        self.position_embedding = PositionEmbedding(config=config)
+        self.encoder = Encoder(config=config)
+        self.decoder = Decoder(config=config)
+        self.embedding = nn.Embedding(config.vocab_size, config.embed_dim)
+        self.head = nn.Linear(config.embed_dim, config.vocab_size)
+        self.drop = nn.Dropout(config.drop_out)
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        tgt: torch.Tensor,
+        src_pad_mask=None,
+        tgt_pad_mask=None,
+    ):
+
+        ## encoder
+
+        src_embedding = self.embedding(src)
+        src_embedding = self.position_embedding(src_embedding)
+        memory = self.encoder.forward(src_embedding, src_pad_mask)
+
+        tgt_embedding = self.embedding(tgt)
+        tgt_embedding = self.position_embedding(tgt_embedding)
+
+        output = self.decoder.forward(tgt_embedding, memory, src_pad_mask, tgt_pad_mask)
+
+        output = self.drop(output)
+
+        output = self.head(output)
+
+        return output

src/sample.py

@@ -0,0 +1,67 @@
+import argparse
+from tokenizers import Tokenizer
+import torch
+from src.config import Config
+from src.model import TranslateModel
+
+
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt_path", default="checkpoints/translate-step=290000.ckpt")
+    parser.add_argument("--zh", default="早上好")
+    return parser.parse_args()
+
+class Inference:
+    def __init__(self,config:Config, ckpt_path):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tokenizer:Tokenizer = Tokenizer.from_file(config.tokenizer_file)
+        self.model:TranslateModel = TranslateModel(config)
+        ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"]
+        state_dict = {}
+        for k, v in ckpt.items():
+            new_k = k[len("net._orig_mod."):]
+            state_dict[new_k] = v
+        self.model.load_state_dict(state_dict, strict=True)
+        self.model.eval()
+        self.model = self.model.to(self.device)
+        self.config = config
+        
+    
+    @torch.no_grad()
+    def sampler(self, src:str)->str:
+        src = self.tokenizer.encode(src).ids
+        tgt = [self.tokenizer.token_to_id("[SOS]")]
+        max_len = self.config.max_len
+        EOS = self.tokenizer.token_to_id("[EOS]")
+
+        src = torch.tensor(src, dtype=torch.long).to(self.device).unsqueeze(0)
+        tgt = torch.tensor(tgt, dtype=torch.long).to(self.device).unsqueeze(0)
+
+        for _ in range(1, max_len):
+            logits = self.model.forward(src, tgt) # [1, len, vocab]
+            logits = logits[:,-1,:]
+            logits = torch.softmax(logits, dim=-1)
+            index = torch.argmax(logits, dim=-1)
+            tgt = torch.cat((tgt, index.unsqueeze(0)), dim=-1)
+            if index.detach().cpu().item() == EOS:
+                break
+        
+        tgt = tgt.detach().cpu().squeeze(0).tolist()
+        tgt_str = self.tokenizer.decode(tgt)
+        return tgt_str
+
+
+def main():
+    args = get_args()
+    config = Config()
+    inference = Inference(config, args.ckpt_path)
+    zh = args.zh
+    result = inference.sampler(zh)
+    print(f"中文:{zh}")
+    print(f"English:{result}")
+
+
+if __name__ == "__main__":
+    main()

src/train.py

@@ -0,0 +1,298 @@
+import os
+import torch
+import argparse
+from torch import nn
+from torch.utils.data import DataLoader, random_split
+from torchmetrics import MeanMetric, MaxMetric
+import lightning as L
+from torchmetrics.classification.accuracy import Accuracy
+from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
+from tokenizers import Tokenizer
+from src.config import Config
+from src.model import TranslateModel
+from src.dataset import TranslateDataset
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning.pytorch.callbacks import RichProgressBar
+from lightning.pytorch.callbacks import ModelCheckpoint
+
+
+import argparse
+
+
+def parser_args():
+    parser = argparse.ArgumentParser(description="Training configuration")
+
+    parser.add_argument("--encoder_layer", type=int, default=6, help="Number of encoder layers")
+    parser.add_argument("--decoder_layer", type=int, default=6, help="Number of decoder layers")
+    parser.add_argument("--embed_dim", type=int, default=512, help="Embedding dimension size")
+    parser.add_argument("--num_heads", type=int, default=8, help="Number of attention heads")
+    parser.add_argument("--drop_out", type=float, default=0.1, help="Dropout rate")
+    parser.add_argument("--max_len", type=int, default=256, help="Maximum sequence length")
+    parser.add_argument("--vocab_size", type=int, default=30000, help="Vocabulary size")
+
+    parser.add_argument("--wmt_zh_en_path", type=str,
+                        default="data/wmt_zh_en_training_corpus.csv",
+                        help="Path to WMT zh-en training corpus")
+    parser.add_argument("--tokenizer_file", type=str,
+                        default="checkpoints/tokenizer.json",
+                        help="Path to tokenizer file")
+    parser.add_argument("--batch_size", type=int, default=64, help="Batch size for training")
+    parser.add_argument("--compile", action="store_true", help="Enable torch.compile if available")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--val_ratio", type=float, default=0.1, help="Validation data ratio")
+    parser.add_argument("--num_workers", type=int, default=4, help="Number of data loading workers")
+    parser.add_argument("--pin_memory", action="store_true", help="Use pinned memory in dataloader")
+
+    parser.add_argument("--tensorboard_dir", type=str, default="log/tensorboard",
+                        help="Directory for tensorboard logs")
+    parser.add_argument("--checkpoint_dir", type=str, default="log/checkpoint",
+                        help="Directory for saving checkpoints")
+
+    # -------- optimizer / scheduler 参数 --------
+    parser.add_argument("--base_lr", type=float, default=3e-4, help="Base learning rate")
+    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.98), help="AdamW betas")
+    parser.add_argument("--eps", type=float, default=1e-9, help="AdamW epsilon")
+    parser.add_argument("--weight_decay", type=float, default=0.1, help="Weight decay")
+
+    parser.add_argument("--warmup_ratio", type=float, default=0.0005,
+                        help="Warmup ratio (fraction of total steps)")
+    parser.add_argument("--start_factor", type=float, default=1e-3,
+                        help="Linear warmup start factor (relative LR scale)")
+    parser.add_argument("--end_factor", type=float, default=1.0,
+                        help="Linear warmup end factor (relative LR scale)")
+    parser.add_argument("--eta_min", type=float, default=3e-6,
+                        help="Minimum LR in cosine annealing")
+
+    parser.add_argument("--max_epochs", type=int, default=10, help="Number of epochs to train")
+
+    # -------- dataset cache 参数 --------
+    parser.add_argument("--data_cache_dir", type=str, default="data/cache.pickle",
+                        help="Path to cache file for dataset")
+    parser.add_argument("--use_cache", action="store_true",
+                        help="Enable dataset caching with pickle")
+    
+    parser.add_argument("--every_n_train_steps", type=int, default=10000,
+                        help="Save checkpoint every N training steps")
+
+    return parser.parse_args()
+
+
+
+def merge_args_config(config: Config, args):
+    for k, v in vars(args).items():
+        setattr(config, k, v)
+    return config
+
+
+class TranslateLitModule(L.LightningModule):
+    def __init__(self, config: Config):
+        super().__init__()
+
+        tokenizer: Tokenizer = Tokenizer.from_file(config.tokenizer_file)
+        self.pad_id = tokenizer.token_to_id("[PAD]")
+        self.net = TranslateModel(config=config)
+        self.train_loss = MeanMetric()
+        self.train_acc = Accuracy(
+            task="multiclass", num_classes=config.vocab_size, ignore_index=self.pad_id
+        )
+        self.criterion = nn.CrossEntropyLoss(ignore_index=self.pad_id)
+
+        self.val_loss = MeanMetric()
+        self.val_acc = Accuracy(
+            task="multiclass", num_classes=config.vocab_size, ignore_index=self.pad_id
+        )
+
+        self.test_loss = MeanMetric()
+        self.test_acc = Accuracy(
+            task="multiclass", num_classes=config.vocab_size, ignore_index=self.pad_id
+        )
+
+        self.val_acc_best = MaxMetric()
+
+        self.config = config
+
+    def forward(self, batch) -> torch.Tensor:
+        pred = self.net.forward(
+            src=batch["src"],
+            tgt=batch["tgt"],
+            src_pad_mask=batch["src_pad_mask"],
+            tgt_pad_mask=batch["tgt_pad_mask"],
+        )
+        return pred
+
+    def on_train_start(self) -> None:
+        self.train_loss.reset()
+        self.train_acc.reset()
+
+        self.val_loss.reset()
+        self.val_acc.reset()
+        self.val_acc_best.reset()
+
+        self.test_loss.reset()
+        self.test_acc.reset()
+
+    def model_step(self, batch):
+
+        logits = self.forward(batch)
+
+        B, L, C = logits.shape
+
+        loss = self.criterion(logits.reshape(-1, C), batch["label"].reshape(-1))
+        preds = torch.argmax(logits, dim=-1)
+        return loss, preds.reshape(-1), batch["label"].reshape(-1)
+
+    def training_step(self, batch, batch_idx):
+        loss, preds, targets = self.model_step(batch)
+
+        # update and log metrics
+        self.train_loss(loss)
+        self.train_acc(preds, targets)
+        self.log(
+            "train/loss", self.train_loss, on_step=True, on_epoch=True, prog_bar=True
+        )
+        self.log(
+            "train/acc", self.train_acc, on_step=True, on_epoch=True, prog_bar=True
+        )
+
+        # return loss or backpropagation will fail
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        pass
+
+    def validation_step(self, batch, batch_idx: int) -> None:
+        loss, preds, targets = self.model_step(batch)
+
+        # update and log metrics
+        self.val_loss(loss)
+        self.val_acc(preds, targets)
+        self.log("val/loss", self.val_loss, on_step=True, on_epoch=True, prog_bar=True)
+        self.log("val/acc", self.val_acc, on_step=True, on_epoch=True, prog_bar=True)
+
+    def on_validation_epoch_end(self) -> None:
+        acc = self.val_acc.compute()  # get current val acc
+        self.val_acc_best(acc)  # update best so far val acc
+        self.log(
+            "val/acc_best", self.val_acc_best.compute(), sync_dist=True, prog_bar=True
+        )
+
+    def test_step(self, batch, batch_idx: int) -> None:
+        loss, preds, targets = self.model_step(batch)
+
+        # update and log metrics
+        self.test_loss(loss)
+        self.test_acc(preds, targets)
+        self.log(
+            "test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True
+        )
+        self.log("test/acc", self.test_acc, on_step=False, on_epoch=True, prog_bar=True)
+
+    def on_test_epoch_end(self) -> None:
+        """Lightning hook that is called when a test epoch ends."""
+        pass
+
+    def setup(self, stage: str) -> None:
+        if self.config.compile and stage == "fit":
+            self.net = torch.compile(self.net)
+
+    def configure_optimizers(self):
+
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.config.base_lr,
+            betas=self.config.betas,
+            eps=self.config.eps,
+            weight_decay=self.config.weight_decay,
+        )
+
+        
+        total_steps = self.trainer.estimated_stepping_batches
+
+        warmup_steps = max(1, int(self.config.warmup_ratio * total_steps))
+        cosine_steps = max(1, total_steps - warmup_steps)
+
+        warmup = LinearLR(
+            optimizer,
+            start_factor=self.config.start_factor,
+            end_factor=self.config.end_factor,
+            total_iters=warmup_steps,
+        )
+
+
+        cosine = CosineAnnealingLR(
+            optimizer,
+            T_max=cosine_steps,
+            eta_min=self.config.eta_min,
+        )
+
+        scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps])
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",   # 每个 step 调整
+                "frequency": 1
+            }
+        }
+
+
+def prepare_dataloader(dataset, config: Config, shuffle=True):
+    dataloader = DataLoader(
+        dataset,
+        batch_size=config.batch_size,
+        shuffle=shuffle,
+        num_workers=config.num_workers,
+        pin_memory=config.pin_memory,
+    )
+    return dataloader
+
+
+def prepare_dataset(config: Config):
+    full_ds = TranslateDataset(config=config)
+    val_ratio = config.val_ratio
+    val_len = int(len(full_ds) * val_ratio)
+    train_len = len(full_ds) - val_len
+    train_ds, val_ds = random_split(
+        full_ds,
+        [train_len, val_len],
+        generator=torch.Generator().manual_seed(config.seed),
+    )
+    return prepare_dataloader(train_ds, config), prepare_dataloader(
+        val_ds, config, False
+    )
+
+
+def prepare_callback(config: Config):
+    logger = TensorBoardLogger(save_dir=config.tensorboard_dir, name="runs")
+    rich_progress_bar = RichProgressBar()
+    checkpoint = ModelCheckpoint(
+        dirpath=config.checkpoint_dir,
+        filename="translate-{step:05d}",
+        save_weights_only=True,
+        every_n_train_steps=config.every_n_train_steps,
+        save_top_k=-1,
+    )
+    return logger, [rich_progress_bar, checkpoint]
+
+
+def main():
+
+    args = parser_args()
+    config = merge_args_config(Config(), args)
+
+    L.seed_everything(config.seed)
+
+    train_loader, val_loader = prepare_dataset(config)
+
+    model = TranslateLitModule(config=config)
+
+    logger, callbacks = prepare_callback(config)
+
+    trainer = L.Trainer(callbacks=callbacks, logger=logger, max_epochs=config.max_epochs)
+
+    trainer.fit(model=model, train_dataloaders=train_loader, val_dataloaders=val_loader)
+
+
+if __name__ == "__main__":
+    main()

src/train_tokenizer.py

@@ -0,0 +1,41 @@
+import csv
+from tokenizers import models, Tokenizer, normalizers, pre_tokenizers, decoders, trainers, processors
+
+
+def text_iterator(file_path):
+    with open(file_path, "r", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            text = row['0'] + " " + row['1']
+            yield text
+
+tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))
+
+tokenizer.normalizer = normalizers.Sequence([
+    normalizers.NFKC()
+])
+
+tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
+
+tokenizer.decoder = decoders.ByteLevel()
+
+trainer = trainers.BpeTrainer(
+    vocab_size=30_000,
+    min_frequency=2,
+    special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"],
+)
+
+tokenizer.train_from_iterator(text_iterator("data/wmt_zh_en_training_corpus.csv"), trainer=trainer)
+
+
+
+tokenizer.post_processor = processors.TemplateProcessing(
+    single="[SOS] $A [EOS]",
+    pair="[SOS] $A [EOS] $B [EOS]",
+    special_tokens=[
+        ("[SOS]", tokenizer.token_to_id("[SOS]")),
+        ("[EOS]", tokenizer.token_to_id("[EOS]")),
+    ],
+)
+
+tokenizer.save("checkpoints/tokenizer.json")