fmeval ライブラリを使用してワークフローをカスタマイズする

カスタム ModelRunner を作成する

1. モデル設定を定義します。

   次のサンプルコードは、Hugging Face モデルのモデル設定を定義できるように、dataclass デコレータをカスタム HFModelConfig クラスに適用する方法を説明しています。

   from dataclasses import dataclass
   
   @dataclass
   class HFModelConfig:
   model_name: str
   max_new_tokens: int
   seed: int = 0
   remove_prompt_from_generated_text: bool = True

   上記のサンプルコードでは、以下が適用されます。

   • パラメータ max_new_tokens は、LLM が返すトークンの数を制限することで、応答の長さを制限するために使用されます。モデルのタイプは、クラスがインスタンス化される際に model_name の値を渡すことで設定されます。この例では、このセクションの最後に示されるとおり、モデル名は gpt2 に設定されています。パラメータ max_new_tokens は、事前トレーニング済みの OpenAI GPT モデルの gpt2 モデル設定を使用してテキスト生成戦略を設定するオプションの 1 つです。その他のモデルタイプについては、「AutoConfig」を参照してください。

   • パラメータ remove_prompt_from_generated_text が True に設定されている場合、生成された応答には、リクエストで送信された元のプロンプトは含まれません。

   その他のテキスト生成パラメータについては、「Hugging Face GenerationConfig のドキュメント」を参照してください。

2. カスタム ModelRunner を作成して、予測メソッドを実装します。次のサンプルコードは、上記のサンプルコードで作成した HFModelConfig クラスを使用して、Hugging Face モデルのカスタム ModelRunner を作成する方法を示しています。

   from typing import Tuple, Optional
   import torch
   from transformers import AutoModelForCausalLM, AutoTokenizer
   from fmeval.model_runners.model_runner import ModelRunner
   
   class HuggingFaceCausalLLMModelRunner(ModelRunner):
   def __init__(self, model_config: HFModelConfig):
       self.config = model_config
       self.model = AutoModelForCausalLM.from_pretrained(self.config.model_name)
       self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
   
   def predict(self, prompt: str) -> Tuple[Optional[str], Optional[float]]:
       input_ids = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
       generations = self.model.generate(
           **input_ids,
           max_new_tokens=self.config.max_new_tokens,
           pad_token_id=self.tokenizer.eos_token_id,
       )
       generation_contains_input = (
           input_ids["input_ids"][0] == generations[0][: input_ids["input_ids"].shape[1]]
       ).all()
       if self.config.remove_prompt_from_generated_text and not generation_contains_input:
           warnings.warn(
               "Your model does not return the prompt as part of its generations. "
               "`remove_prompt_from_generated_text` does nothing."
           )
       if self.config.remove_prompt_from_generated_text and generation_contains_input:
           output = self.tokenizer.batch_decode(generations[:, input_ids["input_ids"].shape[1] :])[0]
       else:
           output = self.tokenizer.batch_decode(generations, skip_special_tokens=True)[0]
   
       with torch.inference_mode():
           input_ids = self.tokenizer(self.tokenizer.bos_token + prompt, return_tensors="pt")["input_ids"]
           model_output = self.model(input_ids, labels=input_ids)
           probability = -model_output[0].item()
   
       return output, probability

   上記のコードでは、FMEval HuggingFaceCausalLLMModelRunner クラスからプロパティを継承するカスタム ModelRunner クラスを使用しています。このカスタムクラスには、コンストラクタと、Tuple を返す予測関数の定義が含まれています。

   その他の ModelRunner の例については、「fmeval ライブラリ」の「model_runner」セクションを参照してください。

   HuggingFaceCausalLLMModelRunner コンストラクタには、以下の定義が含まれています。

   • この設定は、このセクションの冒頭で定義されている HFModelConfig に設定されています。

   • モデルは、インスタンス化時に model_name パラメータを使用して指定された Hugging Face Auto Class の事前トレーニング済みモデルに設定されています。

   • トークナイザーは、model_name で指定された事前トレーニング済みモデルに一致するHugging Faceトークナイザーライブラリのクラスに設定されています。

   HuggingFaceCausalLLMModelRunner クラスの predict メソッドは、以下の定義を使用します。

   • input_ids – モデルの入力を含む変数。モデルは、次のとおり入力を生成します。

     • tokenizer は、prompt に含まれるリクエストをトークン識別子 (ID) に変換します。これらのトークン ID は、特定のトークン (単語、サブ単語、または文字) を表す数値であり、モデルが入力として直接使用できます。トークン ID は、PyTorch が指定したとおり、return_tensors="pt" テンソルオブジェクトとして返されます。その他のタイプの戻りテンソルタイプについては、「apply_chat_template」の「Hugging Face ドキュメント」を参照してください。

     • トークン ID は、モデルが配置されているデバイスに送信され、モデルで使用できるようになります。

   • generations – LLM が生成した応答を含む変数。このモデルの生成関数は、以下の入力を使用して応答を生成します。

     • 前のステップからの input_ids

     • HFModelConfig で指定したパラメータ max_new_tokens

     • pad_token_id は、文末 (eos) トークンを応答に追加します。使用できるその他のトークンについては、「PreTrainedTokenizer」の「Hugging Face ドキュメント」を参照してください。

   • generation_contains_input – 生成された応答に入力プロンプトが含​​まれている場合は True を返し、含まれていない場合は False を返すブール変数。戻り値は、以下の要素単位の比較を使用して計算されます。

     • input_ids["input_ids"][0] に含まれる入力プロンプトのすべてのトークン ID

     • generations[0][: input_ids["input_ids"].shape[1]] に含まれる生成されたコンテンツの先頭

     設定で LLM に remove_prompt_from_generated_text を指示しても、生成された応答に入力プロンプトが含まれていない場合、predict メソッドは警告を返します。

     predict メソッドからの出力には、応答で返されたトークン ID を人間が読み取れるテキストに変換する、batch_decode メソッドが返す文字列が含まれています。remove_prompt_from_generated_text を True に指定した場合、入力プロンプトは生成されたテキストから削除されます。remove_prompt_from_generated_text を False に指定した場合、skip_special_tokens=True で指定されたとおり、ディクショナリ special_token_dict に含めた特別なトークンなしで生成されたテキストが返されます。

3. ModelRunner をテストします。サンプルリクエストをモデルに送信します。

   次の例は、Hugging Face AutoConfig クラスから事前にトレーニングされた gpt2 モデルを使用してモデルをテストする方法を説明しています。

   hf_config = HFModelConfig(model_name="gpt2", max_new_tokens=32)
   model = HuggingFaceCausalLLMModelRunner(model_config=hf_config)

   前のコード例の model_name では、事前トレーニング済みモデルの名前を指定します。HFModelConfig クラスは、パラメータ max_new_tokens の値を使用して hf_config としてインスタンス化され、ModelRunner の初期化に使用されます。

   Hugging Face から別の事前トレーニング済みモデルを使用する場合は、AutoClass の下の from_pretrained で pretrained_model_name_or_path を選択します。

   最後に、ModelRunner をテストします。次のサンプルコードに示されるとおり、モデルにサンプルリクエストを送信します。

   model_output = model.predict("London is the capital of?")[0]
   print(model_output)
   eval_algo.evaluate_sample()