Discordのチャットbotでニューラル機械翻訳を試そう 「JoeyNMT」のカスタマイズについても解説：「Python＋PyTorch」と「JoeyNMT」で学ぶニューラル機械翻訳（2）

精度向上により、近年利用が広まっている「ニューラル機械翻訳」。その仕組みを、自分で動かしながら学んでみましょう。第2回はユースケースごとに「JoeyNMT」をカスタマイズする方法や、Discordのチャットbotに組み込む方法を解説します。

[太田 麻裕美，八楽]

　ハイデルベルク大学の博士課程に在籍しながら、八楽という会社で「ヤラクゼン」の開発に携わっている太田です。ヤラクゼンは、AI翻訳から翻訳文の編集、ドキュメントの共有、翻訳会社への発注までを1つにする翻訳プラットフォームです。

　第1回は、機械翻訳フレームワーク「JoeyNMT」の概要、インストール方法、モデルを訓練する方法を紹介しました。今回は、JoeyNMTをカスタマイズする方法を具体的なユースケースを交えながら紹介します。

　JoeyNMTは、他のフレームワークに比べてコードの行数で9〜10分の1、ファイル数でも4〜5分の1（※1）というミニマルな実装が特長で、核となるモジュールはしっかり入っています。機械学習分野における多くのベンチマークでSOTA（State-of-the-Art）に匹敵するベンチマークスコアを出しています。またデバッグ時にstack traceをたどる際、フラットなディレクトリ構造のおかげで迷わずにエラー箇所を探し当てられるのもメリットです。

※1：OpenNMT-py、XNMTとの比較です。詳細は「Joey NMT: A Minimalist NMT Toolkit for Novices」を参照してください。

　それでは、ユースケースごとにJoeyNMTをカスタマイズする方法を見ていきましょう。

JoeyNMTでトークナイザーを変更するには

　JoeyNMTはデフォルトで「subword-nmt」「sentencepiece」という2つのサブワードトークナイザーに対応しています。では、別のトークナイザーを利用したい場合はどうすればよいでしょうか。

　トークナイザーは「joeynmt/tokenizers.py」で定義できます。例として、「fastBPE」を新しく導入してみましょう。

　fastBPEはsubword-nmtをc++で実装したライブラリです。「SubwordNMTTokenizer」クラスを継承することにします。

class FaseBPETokenizer(SubwordNMTTokenizer):
  def __init__(self, ...):
    try:
      # fastBPEライブラリをインポート
      import fastBPE
    except ImportError as e:
      logger.error(e)
      raise ImportError from e
    super().__init__(level, lowercase, normalize, [...], **kwargs)
    assert self.level == "bpe"
    # codes_path を取得
    self.codes: Path = Path(kwargs["codes_path"])
    assert self.codes.is_file(), f"codes file {self.codes} not found."
    # fastBPEオブジェクト
    self.bpe = fastBPE.fastBPE(self.codes)
    def __call__(self, raw_input: str, is_train: bool = False) -> List[str]:
      #　fastBPE.apply()
      tokenized = self.bpe.apply([raw_input])
      tokenized = tokenized[0].strip().split()
      # 系列の長さが指定の範囲内におさまっているか確認
      if is_train and self._filter_by_length(len(tokenized)):
        return None
      return tokenized

　これでfastBPEでのトークナイズができるようになりました。設定ファイルで「tokenizer_type: "fastbpe"」と選択できるようにするため「_build_tokenizer()」で「FaseBPETokenizer」を呼び出せるようにします。

def _build_tokenizer(cfg: Dict) -> BasicTokenizer:
  [...]
  if tokenizer_type == "sentencepiece": [...]
  elif tokenizer_type == "subword-nmt": [...]
  elif tokenizer_type == "fastbpe":
    assert "codes_path" in tokenizer_cfg
    tokenizer = FaseBPETokenizer(
      level=cfg["level"],
      lowercase=cfg.get("lowercase", False),
      normalize=cfg.get("normalize", False),
      max_length=cfg.get("max_length", -1),
      min_length=cfg.get("min_length", -1),
      **tokenizer_cfg,
    )

　fastBPEにはcodesファイルが必要ですので「codes_path」が設定ファイルで指定されていることを確認しましょう。今回導入した「FaseBPETokenizer」オブジェクトを返すようにしています。

補足

トークナイザーの「__call__()」は、データセットからインスタンスを取り出す際に呼び出されます。例えば「PlaintextDataset」では、「get_item()」内で呼び出されています。

def get_item(self, idx: int, lang: str, is_train: bool = None):
  [...]
  item = self.tokenizer[lang](line, is_train=is_train)
  return item

　つまり、訓練、予測時の「for batch in data_iterator:」のイテレーションで「__getitem__()」がコールされるたびにトークナイズの関数も呼び出されることになります。これは、BPE dropoutを可能にするための実装です。もし、新しく導入するトークナイザーが重い計算を必要としたり、いつも決まった値を返したりするのであれば、データ読み込み時に呼び出される「pre_process()」でトークナイズすることを検討してください（「BaseTokenizer」にある「MosesTokenizer」を利用した事前分割の実装が参考になります）。

JoeyNMTで学習率スケジューラーを変更するには

　JoeyNMTは「torch.optim.lr_scheduler」に入っている「ReduceLROnPlateau」「StepLR」「ExponentialLR」の他、transformerでよく使われる「noamスケジューラー」を実装しています。別の学習率スケジューラーを使いたい場合はどうしたらよいでしょうか？

　学習率スケジューラーは「joeynmt/builders.py」で定義できます。例として、Inverse Square Rootスケジュールを導入してみます。

class BaseScheduler:
  def step(self, step):
    """学習率を更新"""
    self._step = step + 1
    rate = self._compute_rate()
    for p in self.optimizer.param_groups:
      p["lr"] = rate
    self._rate = rate
  def _compute_rate(self):
    raise NotImplementedError

　「BaseScheduler」クラスに、そのステップでの学習率をオプティマイザのパラメーターに渡す部分が実装されています。学習率を計算する「_compute_rate()」関数をオーバーライドします。

　Inverse Square Rootスケジュールは、ステップ数の二乗根に反比例するように学習率を減衰させます。加えて、warmupの期間は、学習率が線形に増加するようにし、warmupの終わりで与えられた学習率に到達するよう係数（decay_rate）を調節します。

class WarmupInverseSquareRootScheduler(BaseScheduler):
  def __init__(
    self,
    optimizer: torch.optim.Optimizer,
    peak_rate: float = 1.0e-3,
    warmup: int = 10000,
    min_rate: float = 1.0e-5,
  ):
    super().__init__(optimizer)
    self.warmup = warmup
    self.min_rate = min_rate
    self.peak_rate = peak_rate
    self.decay_rate = peak_rate * (warmup ** 0.5)
  def _compute_rate(self):
    if step < self.warmup:
      # 線形に増加
      rate = self._step * self.peak_rate / self.warmup
    else:
      # 2乗のルートに反比例
      rate = self.decay_rate * (self._step ** -0.5)
    return max(rate, self.min_rate)

　今回導入したInverse Square Rootスケジューラーを設定ファイルから選択できるように「build_scheduler()」を変更します。

def build_scheduler():
  [...]
  if scheduler_name == "plateau": [...]
  elif scheduler_name == "decaying": [...]
  elif scheduler_name == "exponential": [...]
  elif scheduler_name == "noam": [...]
  elif scheduler_name == "warmupinversesquareroot":
    scheduler = WarmupInverseSquareRootScheduler(
      optimizer=optimizer,
      peak_rate=config.get("learning_rate", 1.0e-3),
      min_rate=config.get("learning_rate_min", 1.0e-5),
      warmup=config.get("learning_rate_warmup", 10000),
    )
    scheduler_step_at = "step"

補足

　訓練を途中で中断した際、その中断したところから再開できるよう、学習率の変数をチェックポイントに保存しています。スケジューラーで保存すべき変数が異なるため、スケジューラーごとに、どの変数を保存するのかを指定する必要があります。

　Inverse Square Rootスケジューラーの場合、デフォルトで保存されるステップ数とそのステップ時の学習率に加えて「warmup」「decay_rate」「peak_rate」「min_rate」を保存します。

class WarmupInverseSquareRootScheduler(BaseScheduler):
  [...]
  def state_dict(self):
    super().state_dict()
    self._state_dict["warmup"] = self.warmup
    self._state_dict["peak_rate"] = self.peak_rate
    self._state_dict["decay_rate"] = self.decay_rate
    self._state_dict["min_rate"] = self.min_rate
    return self._state_dict
  def load_state_dict(self, state_dict):
    super().load_state_dict(state_dict)
    self.warmup = state_dict["warmup"]
    self.decay_rate = state_dict["decay_rate"]
    self.peak_rate = state_dict["peak_rate"]
    self.min_rate = state_dict["min_rate"]

損失関数のカスタマイズ

　機械翻訳では多くの場合、交差エントロピーが損失関数として使われており、JoeyNMTでもデフォルトになっています。損失関数をカスタマイズしたい場合、どうすればよいでしょうか？

　損失関数は「jorynmt/loss.py」で定義できます。第3回で予定している音声翻訳で必要となる「CTC Loss」と呼ばれる損失関数を、少し先取りしてここで導入してみましょう。既存の「XentLoss」クラスを継承して新しいクラス「XentCTCLoss」を作り、PyTorchで実装されているCTC Lossを呼び出します。

　CTC Lossを計算するには、blankを特殊なトークンとして扱う必要があり、そのblankのためのトークンIDを指定しなければなりません。新しくblankトークンを定義してもよいのですが、今回はBOSトークン「<s>」で代用することにします。

class XentCTCLoss(XentLoss):
  def __init__(self,
    pad_index: int,
    bos_index: int,
    smoothing: float = 0.0,
    zero_infinity: bool = True,
    ctc_weight: float = 0.3
  ):
    super().__init__(pad_index=pad_index, smoothing=smoothing)
    self.bos_index = bos_index
    self.ctc_weight = ctc_weight
    self.ctc = nn.CTCLoss(blank=bos_index, reduction='sum')

　「XentCTCLoss」では、すでにある交差エントロピーとCTCの重み付き和を返すようにします。

class XentCTCLoss(XentLoss):
  def forward(self, log_probs, **kwargs) -> Tuple[Tensor, Tensor, Tensor]:
    # CTC Loss の計算に必要な情報がkwargsに入っていることを確認
    assert "trg" in kwargs
    assert "trg_length" in kwargs
    assert "src_mask" in kwargs
    assert "ctc_log_probs" in kwargs
    # 交差エントロピーを計算できるように変形
    log_probs_flat, targets_flat = self._reshape(log_probs, kwargs["trg"])
    # 交差エントロピーを計算
    xent_loss = self.criterion(log_probs_flat, targets_flat)
    # CTC損失を計算
    ctc_loss = self.ctc(
      kwargs["ctc_log_probs"].transpose(0, 1).contiguous(),
      targets=kwargs["trg"], # (seq_length, batch_size)
      input_lengths=kwargs["src_mask"].squeeze(1).sum(dim=1),
      target_lengths=kwargs["trg_length"]
    )
    # 交差エントロピーとCTCの重み付き和を計算
    total_loss = (1.0 - self.ctc_weight) * xent_loss + self.ctc_weight * ctc_loss
    assert total_loss.item() >= 0.0, "loss has to be non-negative."
    return total_loss, xent_loss, ctc_loss

　損失関数は、モデルの「forward()」で呼ばれます。「joeynmt/model.py」の該当部分を変更し「XentCTCLoss」を呼び出せるようにします。

class Model(nn.Module):
  def forward(self, return_type: str = None, **kwargs):
    [...]
    # 通常のデコーダー出力の他、CTCのためのレイヤーからのデコーダー出力も取得
    out, ctc_out = self._encode_decode(**kwargs)
    # デコーダー出力に対し、log_softmax（各トークンの確率）を計算
    log_probs = F.log_softmax(out, dim=-1)
    # バッチごとに損失を計算
    if isinstance(self.loss_function, XentCTCLoss):
      # CTCレイヤーからの出力についても、log_softmaxを計算
      kwargs["ctc_log_probs"] = F.log_softmax(ctc_out, dim=-1)
      # XentCTCLossのforward()を呼び出す
      total_loss, nll_loss, ctc_loss = self.loss_function(log_probs, **kwargs)
    [...]

　バックプロパゲーションに使われるのは重み付き和である「total_loss」だけですが、それぞれの損失関数の学習曲線をプロットするため、「nll_loss」「ctc_loss」も返すようにしています。

補足

　デコーダー（joeynmt/decoders.py）に、CTCLossの計算のためのレイヤーを追加しました。

class TransformerDecoder(Decoder):
  def __init__(self, ...):
    [...]
    self.ctc_output_layer = nn.Linear(encoder_output_size, vocab_size, bias=False)
  def forward(self, ...):
    [...]
    out = self.output_layer(x)
    ctc_output = self.ctc_output_layer(encoder_output)
    return out, x, att, None, ctc_output
class  Model(nn.Module):
  def _encode_decode(self, ...):
    [...]
    out, x, att, _, ctc_out = self._decode(...)
    return out, ctc_out

トークンペナルティで「翻訳結果の繰り返し」を防ぐ

　機械翻訳の出力結果でよくあるのが、繰り返しです。例えば、配布している英日モデルを用いたwmt20テストセットで、以下のような出力を確認しました。

入力："He begged me, "grandma, let me stay, don't do this to me, don't send me back,"" Hernandez said.

出力：「おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん、おばあちゃん」

　根本的には、なぜモデルがこのような繰り返しに高い確率を与えてしまうのかを考える必要があります。ここではその原因には踏み込まず、モデルがこのような繰り返しに高い確率を割り振ったとき、その確率を人為的に低くすることで生成させないという対症療法的な方法を考えます。

　JoeyNMTは、貪欲サーチとビームサーチの2種類の探索を実装しています。どちらも1ステップずつ前から順に生成するauto-regressive、つまりそのステップまでに生成された系列prefixを使って次のトークンを予測します。そこで、そのステップまでに生成された系列prefixを調べ、そこにすでに出現したトークンは、次のトークンを予測する際に確率を下げることにします。

　この例文でいえば「おばあちゃん、おばあちゃん」まで生成したところで次のトークンを予測する際、モデルの予測をそのまま真に受けると「おばあちゃん」が最も確率の高いトークンになってしまいます。そこで、すでにこの系列prefixに出現している「おばあちゃん」のトークンの確率を人為的に下げ、生成されないようにブロックしようというわけです。

　「search.py」の「transformer_greedy()」を見てみましょう。

for step in range(max_output_length):
  with torch.no_grad():
    out, _, _, _ = model(
      return_type="decode",
      trg_input=ys, # すでに生成されたprefixを渡す
      encoder_output=encoder_output,
      encoder_hidden=None,
      src_mask=src_mask,
      unroll_steps=None,
      decoder_hidden=None,
      trg_mask=trg_mask,
      return_attention=return_attention,
    )
  out = out[:, -1]  # logits
  # TODO: repetition penalty / ngram blockerをここで適用
  # もっとも確率が高いトークンを採用
  prob, next_word = torch.max(out, dim=1)

　各ステップで最も確率が高いトークンを採用する前に、モデルの出力（out）を操作してそれまでのステップで生成されたトークンの確率を下げるrepetition penaltyを導入します。

def penalize_repetition(tokens, scores, penalty):
  scores = torch.gather(scores, 1, tokens)  
  scores = torch.where(scores < 0, scores * penalty, scores / penalty)
  scores.scatter_(1, tokens, scores)
  return scores

　ここで「penalty」には1より大きい正の値が入ります。例えば「penalty=2」の場合、すでに出現したトークンの確率を2分の1にせよ、という意味です。

　repetition penaltyは、すでに出現した全てのトークンの確率を一律に下げるように働きます。しかし、例えば日本語の助詞「は」などは複数出現する可能性があり、大きなペナルティーを課したくないときもあるでしょう。そこで、すでに出現した系列prefixのNgramを計算し、次に生成するトークンがそのNgramに一致する場合は確率を0にするという方法もあります。

　仮に「['おばあちゃん', '、', 'おばあちゃん', '、']」という系列prefixがあったとします。3gramの繰り返しをブロックする場合「['おばあちゃん', '、', 'おばあちゃん']」と「['、', 'おばあちゃん', '、']」の2つの3gramがすでに出現していることになります。この系列prefixの次に来るトークンが仮に「'おばあちゃん'」だった場合、直前の2トークンと合わせて「['おばあちゃん', '、', 'おばあちゃん']」となってしまい、すでに出現した3gramのうちの1つと一致してしまいます。すでに出現した3gramと一致するようなトークン「'おばあちゃん'」を禁止トークン（banned_batch_tokens）として扱い、生成されないようにその確率を「float("-inf")」で上書きします。

def block_repeat_ngrams(tokens, scores, no_repeat_ngram_size, step, **kwargs):
  hyp_size = tokens.size(0)  
  banned_batch_tokens = [set([]) for _ in range(hyp_size)]
  trg_tokens = tokens.cpu().tolist()  
  check_end_pos = step + 2 - no_repeat_ngram_size  
  offset = no_repeat_ngram_size - 1
  # 禁止トークンがあるか探します
  for hyp_idx in range(hyp_size):  
    if len(trg_tokens[hyp_idx]) > no_repeat_ngram_size:
      ngram_to_check = trg_tokens[hyp_idx][-offset:]
    for i in range(1, check_end_pos):  # ignore BOS
      if ngram_to_check == trg_tokens[hyp_idx][i:i + offset]:
        banned_batch_tokens[hyp_idx].add(trg_tokens[hyp_idx][i + offset])
  # 見つかった禁止トークンのスコアに対し、-infを代入します。
  for i, banned_tokens in enumerate(banned_batch_tokens):  
    scores[i, list(banned_tokens)] = float("-inf")
  return scores

　オブジェクトをいったんCPUに移し、各シークエンス、各トークンを1つずつループしながら禁止トークンを探していることからも明らかなように、ngram blockerを使うと探索にかかる時間が著しく増大します。GPUの並列化によるアドバンテージを損ないたくない場合は、repetition penaltyを使うことを検討してください。

　ここで「repetition_penaly: 2」と設定して、もう一度同じ例文を全く同じ英日モデルからデコードしてみます。

入力："He begged me, "grandma, let me stay, don't do this to me, don't send me back,"" Hernandez said.

出力：「おばあちゃん、泊まらせてもらって、これもしないで、送ってくれないか」とハーナンデスは言いました。

　意図した通り、一度出現したトークンは生成されにくくなっています。

　「no_repeat_ngram_size: 4」と設定してみます。

入力："He begged me, "grandma, let me stay, don't do this to me, don't send me back,"" Hernandez said.

出力：「おばあちゃん、おばあちゃん、これやらない、送ってくれない」と、ヘルナンデスは言いました。

　4gramより長いフレーズの繰り返しをブロックできています。

補足

　上記の説明では「'おばあちゃん'」が1つのトークンであると仮定していました。配布しているモデルではサブワードトークンを使っており、実際はトークンレベルでの生成は以下のようになっています。

 ['「', 'お', 'ば', 'あ', 'ちゃん', '、', 'お', 'ば', 'あ', 'ちゃ', 'ん', '、', 'これ', 'や', 'ら', 'な', 'い', '、', '送', 'って', 'くれ', 'ない', '」', 'と', '、', 'ヘル', 'ナン', 'デ', 'スは', '言', 'いました', '。', '</s>']

　2回目の「['お', 'ば', 'あ', 'ちゃ', 'ん']」が生成されるところでは、1回目の「['お', 'ば', 'あ', 'ちゃん']」の4gramを避けるために「'ちゃん'」のトークンは選ばれませんでしたが、その代わりに「'ちゃ'」という別のトークンが選ばれ、その次のステップで「'ん'」というトークンが最も高い確率を得ました。その結果、各トークンを結合した出力レベルで見ると、あたかも繰り返しているように見えます。BPEによるトークナイゼーションは何通りもあり得るので、Ngram Blockerを使ってもこのような表層レベルでの繰り返しは起こりえます。

JoeyNMTでアテンション（注意度）を可視化するには

　JoeyNMTには、RNN（Recurrent Neural Network）アーキテクチャからエンコーダーとデコーダー間のアテンションをプロットするオプションが用意されています。Transformerアーキテクチャでアテンションをプロットするにはどうすればよいでしょうか？

　マルチヘッドトランスフォーマーでは、アテンションは1つではありません。全てのレイヤー、全てのヘッドが注意機構で構成されています。エンコーダー層は自己アテンションを、デコーダー層は自己アテンションとクロスアテンションを持っています。今回は、最終レイヤーのクロスアテンションを取り出し、全てのヘッドの平均を取ったものをプロットすることにします。

　マルチヘッドアテンションは「joeynmt/transformer_layers.py」で定義されています。「softmax」を取った後の値を、全てのヘッドで平均して返すようにします。

class MultiHeadedAttention(nn.Module):
  def forward(self, ..., return_weights=False):
    [...]
    attention_weights = self.softmax(scores)
    [...]
    if return_weights:
        # すべてのヘッドの平均を取る [batch_size, query_len, key_len]
      attention_output_weights = attention_weights.view(
        batch_size, self.num_heads, query_len, key_len
      )
      avg_att = attention_output_weights.sum(dim=1) / self.num_heads
      return output, avg_att
    return output, None

　「TransformerDecoderLayer」でクロスアテンションを計算する際に「return_weights」フラグを使えるようにします。

class TransformerDecoderLayer(nn.Module):
  def forward(self, ..., return_attention=False):
    [...]
    h2, att = self.src_trg_att(
      memory, memory, h1, mask=src_mask,
      return_weights=return_attention
    )
    [...]
    out = self.feed_forward(h2)  
    if return_attention:  
      return out, att
    return out, None

　「joeynmt/decoders.py」のトランスフォーマーデコーダーで、最終層のときに「return_attention」フラグをTrueにしてアテンションの重みを取得するようにします。

class TransformerDecoder(nn.Module):
  def forward(self, ...):
    [...]
    last_layer = len(self.layers) - 1  
    for i, layer in enumerate(self.layers):  
      x, att = layer(
        x=x,
        memory=encoder_output,
        src_mask=src_mask,
        trg_mask=trg_mask,
        return_attention=(i == last_layer)
      )
    [...]
    return out, x, att, None

　「joeynmt/search.py」の「transformer_greedy()」で、デコーダーから返ってきたアテンションの値を各ステップでリストに格納し、整形して出力します。

def transformer_greedy(...):
  [...]
  # アテンションの値を取得するかどうか
  return_attention: bool = kwargs.get("return_attention", False)
  [...]
  # アテンションの値を格納するためのプレースホルダー  
  yt = ys.new_zeros((batch_size, 1, src_len), dtype=torch.float)
  [...]
  for step in range(max_output_length):
    # モデルに次のトークンの確率分布を予測させる
    with torch.no_grad():
      out, _, att, _ = model(
        return_type="decode",
        trg_input=ys,
        encoder_output=encoder_output,
        encoder_hidden=None,
        src_mask=src_mask,
        unroll_steps=None,
        decoder_hidden=None,
        trg_mask=trg_mask,
        return_attention=return_attention,
      )
    [...]
    if return_attention:  
      # このステップでデコードした系列prefixの最後のトークンのアテンションの値を格納
      att = att.data[:, -1, :].unsqueeze(1)
      yt = torch.cat([yt, att], dim=1)  # (batch_size, trg_len, src_len)
  [...]
  # BOS-symbol をカット
  output = ys[:, 1:].detach().cpu().numpy()
  attention = yt[:, 1:, :].detach().cpu().numpy() if return_attention else None  
  return output, attention

　これで、トランスフォーマーでもアテンションをプロットできるようになりました。JoeyNMTをテストモードで起動してみましょう。この時に、「--save_attention」オプションを付けると上記の「transformer_greedy()」の「kwargs」に「save_attention=True」が渡されます。

$ python -m joeynmt test config.yaml --save_attention

　テストセットに入っている全ての文のアテンションをプロットしますので、テストセットにはプロットしたい文だけを入れておくようにしてください。

　残念ながら、配布されている日英、英日のモデルではあまりきれいな単語間アラインメントは見られませんでした。アテンションから意味のある単語アラインメントを取り出したい場合は、アラインメントのためのレイヤーを入れるなどの工夫が必要かもしれません。ある程度成功している例として、参考までに独英モデルからプロットしたアテンションもお見せしたいと思います（こちらで配布されています）

補足

　matplotlibの環境によっては、日本語のフォントが文字化けしてしまうかもしれません。その場合は、日本語に対応したフォントを設定する必要があります。上記のプロットにはIPAexGothicを使用しています。「joeynmt/plotting.py」を次のように書き換えてください。

[...]
matplotlib.use("Agg")  
matplotlib.font_manager.fontManager.addfont("/path/to/ipaexg.ttf")
def plot_heatmap(...):
  [...]
  # font config  
  rcParams["xtick.labelsize"] = labelsize  
  rcParams["ytick.labelsize"] = labelsize  
  rcParams['font.family'] = "IPAexGothic"  # support CJK
  [...]

コミュニケーションツール「Discord」用のチャットbotを作ってみよう

　ここからはJoeyNMTで訓練したモデルを、コミュニケーションツール「Discord」上のチャットbotとして動かす方法を紹介します。

Discord アカウントとサーバの準備

　Discord のアカウントがない場合は登録ページでアカウントを作成します。

　続いてサーバを作成するポップアップが開きますので「オリジナルの作成」に進みます。

　「JoeyNMT」という名前のサーバを作ることにします。

　サーバを作成できました。

Bot Applicationの作成

　次に、Discordの開発者ポータルにアクセスし、アプリケーションを新規作成します。

　アプリケーション名を指定します。

　チャットbot追加のボタンをクリックすると確認のポップアップが開くので許可します。

　ここでチャットbot用のアクセストークンが生成されます。後で必要になりますので控えておきます。

　認証に必要なURLを生成します。スコープのセクションでBotを、パーミッションのセクションでAdministratorを選択します。生成されたURLにブラウザからアクセスします。

　認証のポップアップが開きます。ここで、初めに作成したサーバをドロップダウンから選択します。

　管理者権限を与えることを確認して認証します。

　これで設定は一通り終わりました。

　サーバに戻ると、チャットbotが追加されています。

チャットbot用スクリプトの作成

　チャットbot用のスクリプトとして、discord.pyライブラリを使います。discord.pyはpipコマンドでインストールできます。

$ pip install discord.py

　ではスクリプト（discord_joey.py）を書いていきましょう。

　必要なライブラリをインポートします。設定のパートで作成したチャットbotのアクセストークンをスクリプトにコピーします。

[...]
import discord
[...]
# access token
TOKEN = 'your-access-token-here'

　チャットbotには英日、日英のモデルを利用します（学習済みモデルを配布していますのでご利用ください）。JoeyNMTのインタラクティブモードはsingle GPUもしくはCPUで動作します。

CFG_FILES = {
  'en-ja': './models/jparacrawl_enja/config.yaml',
  'ja-en': './models/jparacrawl_jaen/config.yaml'
}
DEVICE = torch.device("cuda")  # DEVICE = torch.device("cpu")
N_GPU = 1  # N_GPU = 0

　イベントを定義します。`on_ready()`でJoeyNMTの学習済みモデルを読み込み、`on_message()`で翻訳を返すようにします。

client = discord.Client()
@client.event
async def on_ready():
  [...] # モデルの読み込み
@client.event
async def on_message(message):
  [...]  # メッセージが来たら翻訳して返す

　モデルの読み込みは「joeynmt/prediction.py」の「translate()」とほぼ同じ手順で行います。

def load_joeynmt_model(cfg_file):
  [...]
  # ボキャブラリを取得
  src_vocab, trg_vocab = build_vocab(cfg["data"], model_dir=model_dir)
  # モデルを構成
  model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)
  # 保存されたチェックポイントからパラメータを読み込む
  ckpt = resolve_ckpt_path(None, load_model, model_dir)
  model_checkpoint = load_checkpoint(ckpt, device=device)
  model.load_state_dict(model_checkpoint["model_state"])
  if device.type == "cuda":
    model.to(device)

　トークナイザー、入力をstringからidに変換するエンコーダーを構成し、インタラクティブモードのためのstream datasetを作ります。stream datasetは初めは空で、入力が来るとその都度キャッシュを更新します。

def load_joeynmt_model(cfg_file):
  [...]
  src_lang = cfg["data"]["src"]["lang"]
  trg_lang = cfg["data"]["trg"]["lang"]
  # トークナイザー
  tokenizer = build_tokenizer(cfg["data"])
  # エンコーダー
  sequence_encoder = {
    src_lang: partial(src_vocab.sentences_to_ids, bos=False, eos=True),
    trg_lang: None,
  }
  # インタラクティブモードのためのデータセットオブジェクト
  test_data = build_dataset(
    dataset_type="stream",
    path=None,
    src_lang=src_lang,
    trg_lang=trg_lang,
    split="test",
    tokenizer=tokenizer,
    sequence_encoder=sequence_encoder,
  )

　幾つかのデコーディングオプションを、インタラクティブモードに対応するように書き換えます。

def load_joeynmt_model(cfg_file):
  [...]
  test_cfg = cfg["testing"]
  test_cfg["batch_type"] = "sentence"
  test_cfg["batch_size"] = 1
  test_cfg["n_best"] = 1
  test_cfg["return_prob"] = "none"
  test_cfg["return_attention"] = False

　メッセージを翻訳する「translate()」では「joeynmt/prediction.py」の「predict()」を呼び出しています。メッセージには、翻訳方向を示す言語タグ「/ja-en/」または「/en-ja/」がついているものとし「get_language_tag()」でこの言語タグと本文を分けています。言語タグの設定に合わせて、翻訳結果を取得します。

@client.event
async def on_message(message):
  # メッセージを言語タグと本文に分ける
  src_input = message.content.strip()
  lang_tag, src_input = get_language_tag(src_input)
  if lang_tag in CFG_FILES:
    # 翻訳を取得
    translation = translate(
      src_input,
      model_dict[lang_tag],
      data_dict[lang_tag],
      cfg_dict[lang_tag],
    )
    # 翻訳結果を返す
    await message.channel.send(translation)

　GitHubのリポジトリに「discord_joey.py」をアップロードしてありますので参考にしてください。では、実行してみます。

$ python discord_joey.py
logged in.
Joey NMT: en-ja model loaded successfully.
Joey NMT: ja-en model loaded successfully.

　モデルがロードされたことを確認したら、Discord上でチャットbotに話し掛けてみます。言語タグを付けるのを忘れずに。

　翻訳結果を返してくれています！　実行されていることが確認できました。

さいごに

　今回はユースケースに合わせてJoeyNMTをカスタマイズする方法を解説しました。同様のシナリオを別のツールキットで実現しようとすると、この何倍ものコードを書き換える必要があります。JoeyNMTの場合、実行スピードを上げるための最適化などはほとんどされておらず、あまり高度なことはできないと感じられた方もいらっしゃるかもしれません。しかし、頭の中で思い描いている変更を愚直に実装できるのはとても大きなアドバンテージだと感じています。

　機械翻訳を良くするアイデアはあっても、既存のフレームワークでは実装が難し過ぎると感じる方、pythonプログラミングや自然言語処理に取り組み始めて日が浅い初心者の方が、JoeyNMTを使って学ぶきっかけになれば幸いです。

　次回は、音声入力からテキスト（文字起こし、翻訳）を生成できるように、JoeyNMTを変更する手順を解説します。