はじめに

最近は複数の大規模言語モデルを組み合わせて使用するシステム（Mixture of experts: MoE, Branch-Train-Merge: BTM）にハマっています。
特にBTMは、独立にモデルを訓練することができるので、「文学専用のAI」、「科学専用のAI」、…のように、専門家の集まりを、わかりやすく構築・統合可能です。
もちろんBTMにも、モデルの構築法、最後の統合方法などの諸課題があります。

今回は、モデルを統合するための簡単な実装コードを書いてみます。
最近は､普通にmergekitもあるようですが､勉強も兼ねた実装です｡

Perplexityとは

データセットのクリーニングにも使えます｡

アプローチ

与えられた入力文章に対するPerplexity（困惑さ）を指標に、使用するモデルを切り替えるシステムを作ります。

イメージ的には、

「文学を学んだモデルを作る」、「科学を学んだモデルを作る」
→「文学系の入力テキストを与える」
→「文学モデルは、文章に馴染み深い（Perplexityが小さい）」、「科学モデルは、文章に馴染みが薄い（Perplexityが大きい）」
→「文学モデルを用いる」

という流れでモデル選択が進みます。

BTMの初期モデルも、このアルゴリズムが使われている印象です。

実装例

モデルの事前訓練をする余裕がないので、今回は試しに、英語が得意なLLama2-7bと、日本語でファインチューニングしたElyza-7bを統合（merge）したシステムを作ってみようと思います。

英語の質問にはllama、日本語の質問にはelyzaで答えることができればコンセプト実証に成功です。

コード

関数とモデルの定義

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer,pipeline
import numpy as np

def perplexity(model, tokenizer, text) -> torch.Tensor:
    tokenized_input = tokenizer.encode(
        text, add_special_tokens=False, return_tensors="pt"
    ).to(model.device)
    with torch.inference_mode():
        output = model(tokenized_input, labels=tokenized_input)
    ppl = torch.exp(output.loss)
    return ppl.item()

class MoE:
    def __init__(self):
        self.models=[]
        self.coef=[]

    def set_coefs(self,coef):
        self.coef=coef

    def append_ELM(self,model,tokenizer):
        pipe=pipeline("text-generation",model=model,tokenizer=tokenizer,
                      max_new_tokens=100
                      )
        self.models.append((model,tokenizer,pipe))
        self.coef.append(1)

    def calc_perplexity(self,text):
        ppl_list=[]
        for model,tokenizer,_ in self.models:
            ppl_list.append(perplexity(model,tokenizer,text))

        return ppl_list

    def ask(self,text,verbose=True):
        ppl_array=np.array(self.calc_perplexity(text))
        ppl_array=ppl_array*np.array(self.coef)
        best_model_id=np.where(ppl_array==min(ppl_array))[0][0]
        if verbose:
            print("perplexity list")
            for i,ppl in enumerate(ppl_array):
                print(i,ppl)
            print(f"model id {best_model_id} is used")
        pipe=self.models[best_model_id][2]
        return pipe(text)[0]['generated_text']

モデルの登録

moe=MoE()


model_name_list =[ 
    "meta-llama/Llama-2-7b-chat-hf",
    "elyza/ELYZA-japanese-Llama-2-7b-instruct",
                  ]

for model_name in model_name_list:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name, device_map="auto",
        torch_dtype=torch.float16
    )

    moe.append_ELM(model,tokenizer)

推論

moe.set_coefs([1,0])

text_list=[
    "hello, I'm John",
    "こんにちは",
    "how can we avoid global warming?",
    "地球温暖化を防ぐにはどうしたらいいか?",
      ]

for text in text_list:
    print("-----")
    print(text)
    response=moe.ask(text)
    print(response)

ポイント

モデルごとにperplexityの大小がかなり違っていたので､今回は便宜的に､coefで補正をかけました。

結果

model 0がllama2 (英語に強いモデル)
model 1がelyza (日本語に強いモデル)
です｡

最後の質疑以外は､狙ったエキスパートモデルが応答してくれました｡

まとめ

perplexityを指標にすることで、与えられた質問ごとに、複数の言語モデルをいい感じにスイッチングできそうなことがわかりました。
スイッチングするアルゴリズム（ルーター）をニューラルネットにしたり、モデル切り替えをtokenレベルで行ってみるなど、諸々の改善はできそうです。

あるいは､既存のキットを使うのも良さそうです｡

3/20追記: Transformersライブラリのpipelineへの埋め込み

transformersライブラリのpipelineから呼び出したかったので、一連の処理を埋め込みました。
かなりのハリボテ作業ですので、ご注意ください(GPT2クラスに偽装しています)。

class

from transformers import GPT2Config, GPT2Model
import torch
import numpy as np


#GPT2クラスを継承します。中身は空です。
class MoEWrapper(GPT2Model):

    config_class = GPT2Config
    #load_tf_weights = load_tf_weights_in_gpt2
    base_model_prefix = "transformer"
    is_parallelizable = True
    supports_gradient_checkpointing = True
    _no_split_modules = ["GPT2Block"]
    _skip_keys_device_placement = "past_key_values"

    verbose=True

    def __init__(self, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)
        self.model_list=[]


    def append_model(self,model):
        self.model_list.append(model)
    
    def set_tokenizer(self,tokenizer):
        self.tokenizer=tokenizer

    def set_model_id(self,model_id):
        self.model=self.model_list[model_id]

    def calc_perplexity(self,tokenized_input):
        ppl_list=[]
        for model in self.model_list:
            ppl_list.append(perplexity(model,tokenized_input))
        return np.array(ppl_list)


    # wrapper functions
　　# generateのノリで、スイッチングする機能を実装すれば、forwardもできるはずです
    #def forward(self,*args, **kwargs):
    #    ret=self.model.forward(*args,**kwargs)
    #    return ret

    def generate(self,input_ids, attention_mask,
                  **generate_kwargs):

        ppl_array=self.calc_perplexity(input_ids)
        best_model_id=np.where(ppl_array==min(ppl_array))[0][0]
        self.set_model_id(best_model_id)
 
        if self.verbose:
            print(f"model {best_model_id} will be used")
            print("ppl array: ",ppl_array)


        ret=self.model.generate(input_ids=input_ids, 
                                attention_mask=attention_mask,
                                  **generate_kwargs)
        return ret


def perplexity(model, tokenized_input) -> torch.Tensor:
    with torch.inference_mode():
        output = model(tokenized_input, labels=tokenized_input)
    ppl = torch.exp(output.loss)
    return ppl.item()

呼び出し

from transformers import GPT2Config, GPT2Model,PreTrainedModel,PretrainedConfig
from MoEWrapper import MoEWrapper
cfg=PretrainedConfig()
cfg=GPT2Config()

moe=MoEWrapper(cfg)

model_name_list =[ 
    "モデル1",
    "モデル2",
                  ]

tokenizer = AutoTokenizer.from_pretrained(model_name_list[0])
for model_name in model_name_list:
    model = AutoModelForCausalLM.from_pretrained(
        model_name, device_map="auto",
        torch_dtype=torch.float16
    )
    moe.append_model(model)

#pipelineで使う
max_new_tokens=100
pipe=pipeline("text-generation",model=moe,
              tokenizer=tokenizer,
              max_new_tokens=max_new_tokens,
                       )

3/26追記 HuggingFaceへの登録