modeling_transliterator.py

import torch
import torch.nn as nn
from model import TransliterationTransformer

class HFTransliterator(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.model = TransliterationTransformer(
            config["src_vocab_size"],
            config["tgt_vocab_size"]
        )

    def forward(self, src, tgt):
        return self.model(src, tgt)