Llama2の日本語化を、デコーダーのカスタムで実現したい！

初めに

Metaが発表した高性能言語モデルであるLlama2

多言語対応して、日本語も対応している、

しかし、返事が英語になってしまう傾向がつよく、日本語での利用で不便。。

Llama2の返答を日本語化する方法として、まず思いつくのは、事後学習として日本語データセットでファインチューニングするという方法である

が、事後学習での問題がある

過学習傾向となりやすく、返答の柔軟性が低下することである（汎用性低下が危惧される）

今回、推論時のデコーダーをカスタムすることで、日本語出力できないかを検討する

方法、考察

Transformerの推論時、デコードをカスタムする
１、日本語のUnicode範囲を指定
２、デコード時の日本語トークンへ有利な調整をする

Llama2などの多言語言語モデルでは、英語が主な学習対象であっても、英語での学習結果が、日本語でも反映されるという多言語横断転移学習という能力が存在する

Llama2は主に英語で学習されているため、日本語で入力しても、英語で出力する傾向が強い

多言語対応しているので、最終出力のニューラルネットワーク全結合層には、日本語での推論結果も出力されてきているのではないかと予想される

普段は英語の出力が目立って見えないが、日本語のみのトークンの範囲内だけでしぼれば、英語と同時に推論された日本語の文章が出力されてきているはず！

そこで、日本語以外のトークン制限を、Unicodeのフィルタリングし、日本語出力を実現する方法を検討した。

まずUnicodeの範囲を以下のように定義した。

# 日本語のUnicode範囲を定義
hiragana_range = (0x3040, 0x309F)
katakana_range = (0x30A0, 0x30FF)
kanji_range = (0x4E00, 0x9FFF)
punctuation_range = (0x3001, 0x3002)  # 「、」と「。」

日本語のトークンリストを作成

# 日本語のトークンのリストを作成
japanese_tokens = [token_id for token_id in range(tokenizer.vocab_size) 
                   if any(char in tokenizer.decode(token_id) for char in tokenizer.decode(token_id) 
                          if ord(char) in range(*hiragana_range) or 
                             ord(char) in range(*katakana_range) or 
                             ord(char) in range(*kanji_range) or
                             ord(char) in range(*punctuation_range))]

プロンプトを設定

# プロンプトの準備
input_text = """<s>[INST]Please tell me the highest mountain in Japan.[/INST]"""
input_ids = tokenizer.encode(input_text, return_tensors="pt")

カスタムデコーダーで推論

def custom_decoder(input_ids, model, tokenizer, max_length=50):
    output_tokens = []

    for _ in range(max_length):
        
        with torch.no_grad():
            outputs = model(input_ids)
        predictions = outputs.logits


        # 日本語でないトークンにペナルティを適用
        non_japanese_tokens = set(range(tokenizer.vocab_size)) - set(japanese_tokens)
        for token_id in non_japanese_tokens:
            predictions[0, -1, token_id] -= 100 

        predicted_token_id = torch.argmax(predictions[0, -1]).item()
        output_tokens.append(predicted_token_id)

        # If the predicted token is a padding token or EOS token, break
        if predicted_token_id in [tokenizer.eos_token_id, tokenizer.pad_token_id]:
            break

        # Add the predicted token to the input for the next prediction
        input_ids = torch.cat([input_ids, torch.tensor([[predicted_token_id]])], dim=-1)


    return tokenizer.decode(output_tokens)

推論の実行

# カスタムデコーダを使用して文章を予測
predicted_sentence = custom_decoder(input_ids, model, tokenizer)
print(f"\nFull predicted sentence: {predicted_sentence}")

今回あえて、英語で質問を入力してみた。
何もしない場合、Llama2の場合、英語での出力される。

今回のカスタムデコーダーの推論結果は以下のようになった

日本語で文章が出力された！

英語で入力しても、英語トークンを隠せば、日本語トークン側でも重みの変化が出力されていることがわかる！

これは多言語横断転移学習の証拠ですね！

まだ内容は少し変ですが。。

現状ではデコード時のUnicode制限だけでは厳しいですね。

次は別の方法も組み合わせます

つづく