Release HuPER Corrector weights and inference code

.gitattributes

@@ -1,35 +1 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

edit_seq_speech/config/vocab.json

@@ -0,0 +1,95 @@
+{
+  "op_to_id": {
+    "KEEP": 0,
+    "DEL": 1,
+    "SUB:<PAD>": 2,
+    "SUB:AA": 3,
+    "SUB:AE": 4,
+    "SUB:AH": 5,
+    "SUB:AO": 6,
+    "SUB:AW": 7,
+    "SUB:AY": 8,
+    "SUB:B": 9,
+    "SUB:CH": 10,
+    "SUB:D": 11,
+    "SUB:DH": 12,
+    "SUB:DX": 13,
+    "SUB:EH": 14,
+    "SUB:ER": 15,
+    "SUB:EY": 16,
+    "SUB:F": 17,
+    "SUB:G": 18,
+    "SUB:HH": 19,
+    "SUB:IH": 20,
+    "SUB:IY": 21,
+    "SUB:JH": 22,
+    "SUB:K": 23,
+    "SUB:L": 24,
+    "SUB:M": 25,
+    "SUB:N": 26,
+    "SUB:NG": 27,
+    "SUB:OW": 28,
+    "SUB:OY": 29,
+    "SUB:P": 30,
+    "SUB:R": 31,
+    "SUB:S": 32,
+    "SUB:SH": 33,
+    "SUB:T": 34,
+    "SUB:TH": 35,
+    "SUB:UH": 36,
+    "SUB:UW": 37,
+    "SUB:V": 38,
+    "SUB:W": 39,
+    "SUB:Y": 40,
+    "SUB:Z": 41,
+    "SUB:ZH": 42
+  },
+  "insert_to_id": {
+    "<NONE>": 0,
+    "<PAD>": 1,
+    "AA": 2,
+    "AE": 3,
+    "AH": 4,
+    "AO": 5,
+    "AW": 6,
+    "AY": 7,
+    "B": 8,
+    "CH": 9,
+    "D": 10,
+    "DH": 11,
+    "DX": 12,
+    "EH": 13,
+    "ER": 14,
+    "EY": 15,
+    "F": 16,
+    "G": 17,
+    "HH": 18,
+    "IH": 19,
+    "IY": 20,
+    "JH": 21,
+    "K": 22,
+    "L": 23,
+    "M": 24,
+    "N": 25,
+    "NG": 26,
+    "OW": 27,
+    "OY": 28,
+    "P": 29,
+    "R": 30,
+    "S": 31,
+    "SH": 32,
+    "T": 33,
+    "TH": 34,
+    "UH": 35,
+    "UW": 36,
+    "V": 37,
+    "W": 38,
+    "Y": 39,
+    "Z": 40,
+    "ZH": 41
+  },
+  "stats": {
+    "num_ops": 43,
+    "num_inserts": 42
+  }
+}

edit_seq_speech/inference.py

@@ -0,0 +1,165 @@
+import torch
+import torch.nn as nn
+import torchaudio
+import json
+import re
+import os
+from g2p_en import G2p
+import pytorch_lightning as pl
+
+from .model import PhonemeCorrector
+from transformers import Wav2Vec2Processor, HubertModel
+
+class PhonemeCorrectionInference:
+    def __init__(self, checkpoint_path, vocab_path, audio_model_name="facebook/hubert-large-ls960-ft", device=None):
+        self.device = device if device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # 1. Load Vocab / Config
+        print(f"Loading config from {vocab_path}...")
+        with open(vocab_path, 'r') as f:
+            self.config = json.load(f)
+            
+        self.op_map = self.config['op_to_id']
+        self.ins_map = self.config['insert_to_id']
+        
+        # Create Reverse Maps (ID -> String)
+        self.id2op = {v: k for k, v in self.op_map.items()}
+        self.id2ins = {v: k for k, v in self.ins_map.items()}
+        
+        # 2. Load G2P
+        self.g2p = G2p()
+        
+        # 3. Load Model
+        print(f"Loading model from {checkpoint_path}...")
+        if os.path.exists(checkpoint_path):
+            checkpoint = torch.load(checkpoint_path, map_location=self.device)
+            hparams = checkpoint.get('hyper_parameters', {})
+            
+            vocab_size = max(self.ins_map.values()) + 1
+            audio_vocab_size = hparams.get('audio_vocab_size', 2048)
+            
+            self.model = PhonemeCorrector.load_from_checkpoint(
+                checkpoint_path,
+                map_location=self.device,
+                vocab_size=vocab_size,
+                audio_vocab_size=audio_vocab_size
+            )
+        else:
+            raise FileNotFoundError(f"Checkpoint not found at {checkpoint_path}")
+            
+        self.model.to(self.device)
+        self.model.eval()
+
+        # 4. Load Audio Tokenizer
+        print(f"Loading Audio Tokenizer: {audio_model_name}")
+        self.audio_processor = Wav2Vec2Processor.from_pretrained(audio_model_name)
+        self.audio_model = HubertModel.from_pretrained(audio_model_name).eval().to(self.device)
+
+    def _clean_phn(self, phn_list):
+        """Standard cleaning to match training."""
+        IGNORED = {"SIL", "'", "SPN", " "} 
+        return [p.rstrip('012') for p in phn_list if p.rstrip('012') not in IGNORED]
+
+    def _get_audio_tokens(self, wav_path):
+        """
+        Runs the audio tokenizer. 
+        IMPORTANT: This must match your training data generation logic.
+        """
+        waveform, sr = torchaudio.load(wav_path)
+        if sr != 16000:
+            resampler = torchaudio.transforms.Resample(sr, 16000)
+            waveform = resampler(waveform)
+            
+        inputs = self.audio_processor(waveform.squeeze().numpy(), return_tensors="pt", sampling_rate=16000)
+        input_values = inputs.input_values.to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.audio_model(input_values)
+            
+        # Placeholder Quantization (Argmax) - Replace if using K-Means
+        features = outputs.last_hidden_state
+        tokens = torch.argmax(features, dim=-1).squeeze()
+        
+        # Downsample to 25Hz (Assuming model is 50Hz)
+        tokens = tokens[::2]
+        return tokens.unsqueeze(0) # (1, T)
+
+    def predict(self, wav_path, text):
+        # A. Prepare Inputs
+        # 1. Text -> Phonemes -> IDs
+        # raw_phns = self.g2p(text)
+        raw_phns = text.split()  # Assuming input text is already phonemized for inference
+        src_phns = self._clean_phn(raw_phns)
+        
+        # Create text vocab from insert_to_id (same as dataset)
+        text_vocab = {k: v for k, v in self.ins_map.items() if k not in ['<NONE>', '<PAD>']}
+        text_ids = [text_vocab.get(p, text_vocab.get("AA", 2)) for p in src_phns]
+        text_tensor = torch.tensor([text_ids], dtype=torch.long).to(self.device)
+        
+        # 2. Audio -> Tokens
+        audio_tensor = self._get_audio_tokens(wav_path)
+        
+        # B. Run Model
+        with torch.no_grad():
+            # Create masks
+            txt_mask = torch.ones_like(text_tensor)
+            aud_mask = torch.ones_like(audio_tensor)
+            
+            logits_op, logits_ins = self.model(
+                text_tensor, audio_tensor, txt_mask, aud_mask
+            )
+            
+            # C. Decode
+            pred_ops = torch.argmax(logits_op, dim=-1).squeeze().tolist()
+            pred_ins = torch.argmax(logits_ins, dim=-1).squeeze().tolist()
+            
+        # Ensure lists
+        if not isinstance(pred_ops, list): pred_ops = [pred_ops]
+        if not isinstance(pred_ins, list): pred_ins = [pred_ins]
+
+        # D. Reconstruct Sequence
+        final_phonemes = []
+        log = []
+        
+        for i, (orig, op_id, ins_id) in enumerate(zip(src_phns, pred_ops, pred_ins)):
+            
+            # 1. Apply Operation
+            op_str = self.id2op.get(op_id, "KEEP")
+            curr_log = {"src": orig, "op": op_str, "ins": "NONE"}
+            
+            if op_str == "KEEP":
+                final_phonemes.append(orig)
+            elif op_str == "DEL":
+                pass # Do not append
+            elif op_str.startswith("SUB:"):
+                # Extract phoneme: "SUB:AA" -> "AA"
+                new_phn = op_str.split(":")[1]
+                final_phonemes.append(new_phn)
+            
+            # 2. Apply Insertion
+            ins_str = self.id2ins.get(ins_id, "<NONE>")
+            if ins_str != "<NONE>":
+                final_phonemes.append(ins_str)
+                curr_log["ins"] = ins_str
+                
+            log.append(curr_log)
+            
+        return final_phonemes, log
+
+if __name__ == "__main__":
+    ckpt_path = "/data/chenxu/checkpoints/edit_seq_speech/phoneme-corrector/last.ckpt"
+    vocab_path = "edit_seq_speech/config/vocab.json"
+    wav_file = "test.wav"
+    text_input = "Last Sunday"
+    
+    if os.path.exists(ckpt_path) and os.path.exists(wav_file):
+        infer = PhonemeCorrectionInference(ckpt_path, vocab_path)
+        result, details = infer.predict(wav_file, text_input)
+        
+        print(f"Input Text: {text_input}")
+        print(f"Result Phn: {result}")
+        print("-" * 20)
+        for step in details:
+            print(f"{step['src']} -> {step['op']} + Insert({step['ins']})")
+    else:
+        print("Please set valid paths for checkpoint and wav file.")

edit_seq_speech/model.py

@@ -0,0 +1,330 @@
+import torch
+import torch.nn as nn
+import math
+import pytorch_lightning as pl
+
+class PhonemeCorrector(pl.LightningModule):
+    def __init__(self, vocab_size, audio_vocab_size, d_model=256, nhead=4, num_layers=4, dropout=0.1, lr=1e-4, 
+                 weight_decay=0.01, scheduler_config=None, optimizer_config=None):
+        super().__init__()
+        self.save_hyperparameters()
+        self.scheduler_config = scheduler_config or {}
+        self.optimizer_config = optimizer_config or {}
+        
+        # 1. Embeddings
+        self.text_embedding = nn.Embedding(vocab_size, d_model)
+        self.audio_embedding = nn.Embedding(audio_vocab_size, d_model)
+        
+        # Positional Encoding (Standard Sinusoidal)
+        self.pos_encoder = PositionalEncoding(d_model, dropout)
+        
+        # 2. The Core Transformer (Text querying Audio)
+        decoder_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
+        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
+        
+        # 3. Prediction Heads - 2-head architecture
+        # Head 1: Operation (KEEP, DEL, SUB:AA, SUB:AE, ...)
+        # num_ops = vocab_size + 2 (KEEP=0, DEL=1, SUB:phonemes=2+)
+        # This matches the precomputed op_ids format
+        num_ops = vocab_size + 2
+        self.head_op = nn.Linear(d_model, num_ops)
+        
+        # Head 2: Insertion (NONE=0, AA, AE, ...)
+        # num_inserts = vocab_size (NONE=0, then phonemes)
+        num_inserts = vocab_size
+        self.head_ins = nn.Linear(d_model, num_inserts)
+
+    def forward(self, text_ids, audio_ids, text_mask=None, audio_mask=None):
+        """
+        text_ids: (Batch, Text_Len)
+        audio_ids: (Batch, Audio_Len)
+        masks: (Batch, Len) - 1 for valid, 0 for pad.
+        """
+        text_emb = self.pos_encoder(self.text_embedding(text_ids))
+        audio_emb = self.pos_encoder(self.audio_embedding(audio_ids))
+        
+        txt_pad_mask = (text_mask == 0) if text_mask is not None else None
+        aud_pad_mask = (audio_mask == 0) if audio_mask is not None else None
+
+        encoded_features = self.transformer(
+            tgt=text_emb, 
+            memory=audio_emb,
+            tgt_key_padding_mask=txt_pad_mask,
+            memory_key_padding_mask=aud_pad_mask
+        )
+        
+        logits_op = self.head_op(encoded_features)
+        logits_ins = self.head_ins(encoded_features)
+        
+        return logits_op, logits_ins
+
+    def training_step(self, batch, batch_idx):
+        input_ids = batch['input_ids']
+        audio_tokens = batch['audio_tokens']
+        lbl_op = batch['labels']['op']
+        lbl_ins = batch['labels']['ins']
+        txt_mask = batch['masks']['text']
+        audio_mask = batch['masks']['audio']
+        
+        logits_op, logits_ins = self(input_ids, audio_tokens, txt_mask, audio_mask)
+        
+        # Active loss mask (only compute loss on valid text tokens)
+        active_loss = txt_mask.view(-1) == 1
+        
+        # OP LOSS (includes KEEP, DEL, and all SUB:phoneme operations)
+        num_ops = self.hparams.vocab_size + 2
+        loss_op = nn.functional.cross_entropy(
+            logits_op.view(-1, num_ops)[active_loss], 
+            lbl_op.view(-1)[active_loss]
+        )
+        
+        # INS LOSS
+        loss_ins = nn.functional.cross_entropy(
+            logits_ins.view(-1, self.hparams.vocab_size)[active_loss],
+            lbl_ins.view(-1)[active_loss]
+        )
+        
+        loss = loss_op + loss_ins
+        self.log('train_loss', loss, prog_bar=True)
+        self.log('train_loss_op', loss_op)
+        self.log('train_loss_ins', loss_ins)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        input_ids = batch['input_ids']
+        audio_tokens = batch['audio_tokens']
+        lbl_op = batch['labels']['op']
+        lbl_ins = batch['labels']['ins']
+        txt_mask = batch['masks']['text']
+        audio_mask = batch['masks']['audio']
+        
+        logits_op, logits_ins = self(input_ids, audio_tokens, txt_mask, audio_mask)
+        
+        # Compute losses
+        active_loss = txt_mask.view(-1) == 1
+        num_ops = self.hparams.vocab_size + 2
+        
+        loss_op = nn.functional.cross_entropy(
+            logits_op.view(-1, num_ops)[active_loss], 
+            lbl_op.view(-1)[active_loss]
+        )
+        
+        loss_ins = nn.functional.cross_entropy(
+            logits_ins.view(-1, self.hparams.vocab_size)[active_loss],
+            lbl_ins.view(-1)[active_loss]
+        )
+        
+        loss = loss_op + loss_ins
+        
+        # Compute accuracy
+        pred_op = torch.argmax(logits_op, dim=-1)
+        pred_ins = torch.argmax(logits_ins, dim=-1)
+        
+        # OP accuracy
+        op_correct = (pred_op == lbl_op) & txt_mask
+        op_acc = op_correct.sum().float() / txt_mask.sum().float()
+        
+        # INS accuracy
+        ins_correct = (pred_ins == lbl_ins) & txt_mask
+        ins_acc = ins_correct.sum().float() / txt_mask.sum().float()
+        
+        # Overall accuracy: correct OP prediction
+        overall_acc = op_acc
+        
+        # Per-operation accuracy (KEEP=0, DEL=1, SUB>=2)
+        keep_mask = (lbl_op == 0) & txt_mask
+        del_mask = (lbl_op == 1) & txt_mask
+        sub_op_mask = (lbl_op >= 2) & txt_mask
+        
+        keep_acc = torch.tensor(0.0, device=loss.device)
+        del_acc = torch.tensor(0.0, device=loss.device)
+        sub_op_acc = torch.tensor(0.0, device=loss.device)
+        
+        if keep_mask.sum() > 0:
+            keep_correct = (pred_op == lbl_op) & keep_mask
+            keep_acc = keep_correct.sum().float() / keep_mask.sum().float()
+        
+        if del_mask.sum() > 0:
+            del_correct = (pred_op == lbl_op) & del_mask
+            del_acc = del_correct.sum().float() / del_mask.sum().float()
+        
+        if sub_op_mask.sum() > 0:
+            sub_op_correct = (pred_op == lbl_op) & sub_op_mask
+            sub_op_acc = sub_op_correct.sum().float() / sub_op_mask.sum().float()
+        
+        # Log metrics
+        self.log('val_loss', loss, prog_bar=True, sync_dist=True)
+        self.log('val_loss_op', loss_op, sync_dist=True)
+        self.log('val_loss_ins', loss_ins, sync_dist=True)
+        self.log('val_acc', overall_acc, prog_bar=True, sync_dist=True)
+        self.log('val_acc_op', op_acc, sync_dist=True)
+        self.log('val_acc_ins', ins_acc, sync_dist=True)
+        self.log('val_acc_keep', keep_acc, sync_dist=True)
+        self.log('val_acc_del', del_acc, sync_dist=True)
+        self.log('val_acc_sub_op', sub_op_acc, sync_dist=True)
+        
+        return {
+            'val_loss': loss,
+            'val_acc': overall_acc,
+            'val_acc_op': op_acc,
+            'val_acc_ins': ins_acc
+        }
+
+    def configure_optimizers(self):
+        # Get optimizer configuration
+        optimizer_name = self.optimizer_config.get("name", "adamw").lower()
+        lr = self.hparams.lr
+        weight_decay = getattr(self.hparams, 'weight_decay', 0.01)
+        
+        if optimizer_name == "adamw":
+            optimizer = torch.optim.AdamW(
+                self.parameters(),
+                lr=lr,
+                weight_decay=weight_decay,
+                betas=self.optimizer_config.get("betas", [0.9, 0.999]),
+                eps=self.optimizer_config.get("eps", 1.0e-8)
+            )
+        elif optimizer_name == "adam":
+            optimizer = torch.optim.Adam(
+                self.parameters(),
+                lr=lr,
+                weight_decay=weight_decay,
+                betas=self.optimizer_config.get("betas", [0.9, 0.999]),
+                eps=self.optimizer_config.get("eps", 1.0e-8)
+            )
+        else:
+            raise ValueError(f"Unknown optimizer: {optimizer_name}")
+        
+        # Configure scheduler
+        scheduler_type = self.scheduler_config.get("type", "cosine").lower()
+        
+        # Calculate total training steps
+        max_epochs = getattr(self.trainer, 'max_epochs', 50)
+        if self.trainer and hasattr(self.trainer, 'estimated_stepping_batches'):
+            total_steps = self.trainer.estimated_stepping_batches
+        else:
+            # Fallback: estimate steps per epoch
+            estimated_steps_per_epoch = 1000  # Conservative estimate
+            total_steps = max_epochs * estimated_steps_per_epoch
+        
+        warmup_ratio = self.scheduler_config.get("warmup_ratio", 0.1)
+        warmup_steps = max(1, int(total_steps * warmup_ratio))
+        
+        if scheduler_type == "cosine":
+            # Use transformers' cosine scheduler with warmup
+            try:
+                from transformers import get_cosine_schedule_with_warmup
+                eta_min = self.scheduler_config.get("eta_min", 1.0e-6)
+                scheduler = get_cosine_schedule_with_warmup(
+                    optimizer,
+                    num_warmup_steps=warmup_steps,
+                    num_training_steps=total_steps,
+                    num_cycles=0.5,  # Default cosine cycles
+                    last_epoch=-1
+                )
+            except ImportError:
+                # Fallback to PyTorch implementation
+                from torch.optim.lr_scheduler import LambdaLR
+                import math
+                eta_min = self.scheduler_config.get("eta_min", 1.0e-6)
+                def lr_lambda(step):
+                    if step < warmup_steps:
+                        return step / warmup_steps
+                    else:
+                        # Cosine annealing after warmup
+                        progress = (step - warmup_steps) / (total_steps - warmup_steps)
+                        cosine_value = 0.5 * (1 + math.cos(math.pi * progress))
+                        return eta_min / lr + (1 - eta_min / lr) * cosine_value
+                scheduler = LambdaLR(optimizer, lr_lambda)
+            
+        elif scheduler_type == "linear":
+            # Use transformers' linear scheduler with warmup
+            try:
+                from transformers import get_linear_schedule_with_warmup
+                scheduler = get_linear_schedule_with_warmup(
+                    optimizer,
+                    num_warmup_steps=warmup_steps,
+                    num_training_steps=total_steps
+                )
+            except ImportError:
+                # Fallback to PyTorch implementation
+                from torch.optim.lr_scheduler import LambdaLR
+                def lr_lambda(step):
+                    if step < warmup_steps:
+                        return step / warmup_steps
+                    else:
+                        progress = (step - warmup_steps) / (total_steps - warmup_steps)
+                        return max(0.0, 1.0 - progress)
+                scheduler = LambdaLR(optimizer, lr_lambda)
+            
+        elif scheduler_type == "polynomial":
+            # Use transformers' polynomial scheduler with warmup
+            try:
+                from transformers import get_polynomial_decay_schedule_with_warmup
+                power = self.scheduler_config.get("power", 1.0)
+                scheduler = get_polynomial_decay_schedule_with_warmup(
+                    optimizer,
+                    num_warmup_steps=warmup_steps,
+                    num_training_steps=total_steps,
+                    power=power
+                )
+            except ImportError:
+                # Fallback: use linear scheduler
+                from torch.optim.lr_scheduler import LambdaLR
+                def lr_lambda(step):
+                    if step < warmup_steps:
+                        return step / warmup_steps
+                    else:
+                        progress = (step - warmup_steps) / (total_steps - warmup_steps)
+                        return max(0.0, (1.0 - progress) ** power)
+                scheduler = LambdaLR(optimizer, lr_lambda)
+            
+        elif scheduler_type == "reduce_on_plateau":
+            from torch.optim.lr_scheduler import ReduceLROnPlateau
+            scheduler = ReduceLROnPlateau(
+                optimizer,
+                mode='min',
+                factor=self.scheduler_config.get("factor", 0.5),
+                patience=self.scheduler_config.get("patience", 3),
+                min_lr=self.scheduler_config.get("min_lr", 1.0e-6),
+                verbose=True
+            )
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "monitor": "train_loss",
+                    "interval": "epoch",
+                    "frequency": 1,
+                }
+            }
+        else:
+            # No scheduler
+            return optimizer
+        
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1,
+            }
+        }
+
+# Helper for Positional Encoding
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe.unsqueeze(0))
+
+    def forward(self, x):
+        # x: (Batch, Seq, Dim)
+        x = x + self.pe[:, :x.size(1)]
+        return self.dropout(x)