⛓

huggingfaceのモデルをlangchainで使う＋量子化とストップワード（停止文字列）の設定

2024/04/01に公開

Python

はじめに

今回は、huggingfaceからモデル（主にbinやsafetensorsの形式）をlangchainで扱う場合の方法と同時に設定しておきたい量子化とストップワードの設定について紹介します。

と、その前に（一応）gguf形式のモデルを扱う方法を先に紹介します。

gguf形式のモデルを扱う場合

gguf_sample.py

from langchain_community.llms.llamacpp import LlamaCpp
from langchain_core.prompts import PromptTemplate
from langchain.chains.llm import LLMChain

# モデルのパス
model_path = "./ELYZA-japanese-Llama-2-7b-fast-instruct-q4_K_M.gguf"

# プロンプトテンプレートの定義
question_prompt_template = """
[INST]
{question} 
[/INST]
"""

# プロンプトの設定
QUESTION_PROMPT = PromptTemplate(
    template=question_prompt_template, # プロンプトテンプレートをセット
    input_variables=["question"] # プロンプトに挿入する変数
)

# モデルの設定
llm = LlamaCpp(
    model_path=model_path, # ダウンロードしたモデルのローカルパス
    n_gpu_layers=25, # gpuに処理させるlayerの数
    # stop=["コンピュータ", "人間"], # 停止文字列
)

# 質問回答chainの設定
chain = LLMChain(llm=llm, prompt=QUESTION_PROMPT)

# 質問文
question = "AIとは何ですか？"

# LLMの回答生成
response = chain.invoke(question)

# 回答を確認
print(response["text"])

gguf形式のものは上記のように（langchain_communityから）llamacppを使うことでchainに渡せます。ですが、binやsafetensorsなどgguf以外の形式はllamacppでは扱えません。

HuggingFacePipelineを使った実行

HuggingFaceのモデルをlangchainで扱うためにHuggingFacePipelineというものが用意されています。

まずは量子化なしの状態で実行を試してみます。
今回は以下のモデル（googleのgemma-2b-it）を使いました。

※量子化なしでの実行が厳しいため小さめのモデルを選びました。

必要なライブラリをインストールしましょう。

$ pip install langchain
$ pip install langchain-community
$ pip install transformers // gemmaはv4.38.1以上必要です。
$ pip install accelerate
$ pip install torch

エラーが出る場合は「pip install -U transformers」のように一度updateすることをおすすめします。

コーデイング

pipeline_sample.py

import os
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline
)
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_core.prompts import PromptTemplate
from langchain.chains.llm import LLMChain

# huggingfaceのトークンの設定（gemmaを使用するのに必要なため）
os.environ["HF_TOKEN"] = "<< huggingfaceのトークン >>"

# gemma-2b-it のリポジトリIDを設定
repo_id = "google/gemma-2b-it"

# モデルの設定
model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    device_map={"": 'cuda:0'} # gpu（cuda）での実行を指定
)

# トークナイザーの設定
tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path=repo_id
)

# パイプラインの設定
pipe = pipeline(
    "text-generation", # タスクの指定
    model=model, # モデルをセット
    tokenizer=tokenizer, # トークナイザーをセット
    max_new_tokens=64 # 生成するトークンの最大数
)

# langchainで扱える形式に変換
llm = HuggingFacePipeline(
    pipeline=pipe # パイプラインをセット
)

# プロンプトテンプレートをフォーマット
question_prompt_template_format = tokenizer.apply_chat_template(
    conversation = [
        {"role": "user", "content": "{question}"}
    ], 
    tokenize=False, 
    add_generation_prompt=True
)

# プロンプトの設定
QUESTION_PROMPT = PromptTemplate(
    template=question_prompt_template_format,
    input_variables=["question"]
)

# 質問回答chainの設定
chain = LLMChain(llm=llm, prompt=QUESTION_PROMPT)

# 質問文 (DeepL翻訳: AIについて50字程度で簡潔に説明してください。）
question = 'Briefly describe AI in 50 words or so.'

# LLMの回答生成
response = chain.invoke({"question": question})

# 回答を確認
print(response["text"])

※上記、huggingfaceトークンは実行しやすくするためハードコーディングしています（本来は切り分けた方がいいです・・・）。

langchainで実行するための変換コード（HuggingFacePipelineの設定）は以下の部分です。

pipeline_sample.py

# パイプラインの設定
pipe = pipeline(
    "text-generation", # タスクの指定
    model=model, # モデルをセット
    tokenizer=tokenizer, # トークナイザーをセット
    max_new_tokens=64 # 生成するトークンの最大数
)

# langchainで扱える形式に変換
llm = HuggingFacePipeline(
    pipeline=pipe # パイプラインをセット
)

上記のようにモデルとトークナイザーを渡してpipelineを設定 → HuggingFacePipelineにpipelineをセットすることでlangchainで実行可能な形式に変換できます。

出力確認

Artificial intelligence (AI) is the simulation of human intelligence processes by machines, enabling them to perform tasks that typically require human intelligence, such as learning, problem-solving, and decision-making.

DeepL翻訳: 人工知能（AI）とは、人間の知能プロセスを機械がシミュレーションすることで、学習、問題解決、意思決定など、通常人間の知能を必要とするタスクを機械が実行できるようにするものである。

HuggingFacePipelineを使った実行（量子化）

今回の量子化にはBitsAndBytesConfigを使います。

必要なライブラリをインストールしましょう。

$ pip install bitsandbytes

コーディング

pipe_4bit_sample.py

import os
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
    BitsAndBytesConfig
)
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_core.prompts import PromptTemplate
from langchain.chains.llm import LLMChain

# huggingfaceのトークンの設定（gemmaを使用するのに必要なため）
os.environ["HF_TOKEN"] = "<< huggingfaceのトークン >>"

# gemma-2b-it のリポジトリIDを設定
repo_id = "google/gemma-2b-it"

# 量子化の構成を設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True, # 4bit量子化を使用
    bnb_4bit_quant_type="nf4", # 量子化データ型にNF4を設定
    bnb_4bit_compute_dtype=torch.bfloat16 # 量子化計算時のデータ型にbfloat16を設定
)

# モデルの設定
model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    device_map={"": 'cuda:0'}, # gpu（cuda）での実行を指定
    quantization_config=quantization_config # 量子化の構成をセット
)

# トークナイザーの設定
tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path=repo_id
)

# パイプラインの設定
pipe = pipeline(
    "text-generation", # タスクの指定
    model=model, # モデルをセット
    tokenizer=tokenizer, # トークナイザーをセット
    max_new_tokens=64 # 生成するトークンの最大数
)

# langchainで扱える形式に変換
llm = HuggingFacePipeline(
    pipeline=pipe # パイプラインをセット
)

# プロンプトテンプレートをフォーマット
question_prompt_template_format = tokenizer.apply_chat_template(
    conversation = [
        {"role": "user", "content": "{question}"}
    ], 
    tokenize=False, 
    add_generation_prompt=True
)

# プロンプトの設定
QUESTION_PROMPT = PromptTemplate(
    template=question_prompt_template_format,
    input_variables=["question"]
)

# 質問回答chainの設定
chain = LLMChain(llm=llm, prompt=QUESTION_PROMPT)

# 質問文（和訳: AIについて50字程度で簡潔に説明してください。）
question = 'Briefly describe AI in 50 words or so.'

# LLMの回答生成
response = chain.invoke({"question": question})

# 回答を確認
print(response["text"])

BitsAndBytesConfigを使った量子化設定は以下の部分です。

pipe_4bit_sample.py

# 量子化の構成を設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True, # 4bit量子化を使用
    bnb_4bit_quant_type="nf4", # 量子化データ型にNF4を設定
    bnb_4bit_compute_dtype=torch.bfloat16 # 量子化計算時のデータ型にbfloat16を設定
)

# モデルの設定
model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    device_map={"": 'cuda:0'}, # gpu（cuda）での実行を指定
    quantization_config=quantization_config # 量子化の構成をセット
)

上記のようにBitsAndBytesConfigで量子化の構成を設定する → モデル設定の際に量子化の構成をセットすることで量子化が行えます。
BitsAndBytesConfigでは load_in_4bit=True → load_in_8bit=True とすることで8bit量子化も選択可能です。

pipe_8bit_sample.py

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True
)

上記以外にも bnb_4bit_use_double_quant（二重量子化）などの量子化手法が設定可能です（速度やメモリ消費の削減などの効果を期待できます)。
設定できるパラメータに関しては以下を参照してください。

出力確認

Artificial intelligence (AI) is the ability of a machine to perform tasks that typically require human intelligence, such as learning, problem-solving, and decision-making.

DeepL翻訳: 人工知能（AI）とは、学習、問題解決、意思決定など、通常は人間の知性を必要とするタスクを機械が実行する能力のことである。

HuggingFacePipelineを使った実行（量子化 + ストップワード）

ストップワードの設定はStoppingCriteriaを使います。

コーディング

pipe_4bit_stop_sample.py

import os
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
    BitsAndBytesConfig,
    StoppingCriteria,
    StoppingCriteriaList
)
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_core.prompts import PromptTemplate
from langchain.chains.llm import LLMChain

# huggingfaceのトークンの設定（gemmaを使用するのに必要なため）
os.environ["HF_TOKEN"] = '<< huggingfaceのトークン >>'

# gemma-2b-it のリポジトリIDを設定
repo_id = "google/gemma-2b-it"

# 量子化の構成を設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True, # 4bit量子化を使用
    bnb_4bit_quant_type="nf4", # 量子化データ型にNF4を設定
    bnb_4bit_compute_dtype=torch.bfloat16 # 量子化計算時のデータ型にbfloat16を設定
)

# モデルの設定
model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    device_map={"": 'cuda:0'}, # gpu（cuda）での実行を指定
    quantization_config=quantization_config # 量子化の構成をセット
)

# トークナイザーの設定
tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    use_fast = False,
)

# 生成トークンの停止条件設定クラス
class MyStoppingCriteria(StoppingCriteria):
    def __init__(self, stop_words_ids=[]):
      super().__init__()
      self.stop_words_ids = stop_words_ids
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
        for stop_word_ids in self.stop_words_ids: 
            # 生成した末尾のトークンがストップワードと一致するか確認
            if input_ids[0][-1] == stop_word_ids: 
                # トークンの生成をストップ
                return True
        return False

# ストップワードのリスト
stop_words = [" machine", " human"]

stop_words_ids = []
for stop_word in stop_words:
    # 文字列をテンソルに変換
    stop_word_ids = tokenizer(stop_word, return_tensors='pt', add_special_tokens=False)['input_ids']
    # テンソルをcudaに配置
    stop_word_ids = stop_word_ids.to("cuda:0") 
    # テンソルをリストに格納
    stop_words_ids.append(stop_word_ids) 

# MyStoppingCriteriaのインスタンスを生成
my_stopping_criteria = MyStoppingCriteria(stop_words_ids=stop_words_ids) 

# StoppingCriteriaListのインスタンスを生成
stopping_criteria_list = StoppingCriteriaList([my_stopping_criteria]) 

# パイプラインの設定
pipe = pipeline(
    "text-generation", # タスクの指定
    model=model, # モデルをセット
    tokenizer=tokenizer, # トークナイザーをセット
    max_new_tokens=64, # 生成するトークンの最大数
    stopping_criteria=stopping_criteria_list # StoppingCriteriaListのインスタンスをセット
)

# langchainで扱える形式に変換
llm = HuggingFacePipeline(
    pipeline=pipe # パイプラインをセット
)

# プロンプトテンプレートをフォーマット
question_prompt_template_format = tokenizer.apply_chat_template(
    conversation = [
        {"role": "user", "content": "{question}"}
    ], 
    tokenize=False, 
    add_generation_prompt=True
)

# プロンプトの設定
QUESTION_PROMPT = PromptTemplate(
    template=question_prompt_template_format,
    input_variables=["question"]
)

# 質問回答chainの設定
chain = LLMChain(llm=llm, prompt=QUESTION_PROMPT)

# 質問文（和訳: AIについて50字程度で簡潔に説明してください。）
question = 'Briefly describe AI in 50 words or so.'

# LLMの回答生成
response = chain.invoke({"question": question})

# 回答を確認
print(response["text"])

StoppingCriteriaを使ったストップワードの設定は以下の部分です。

pipe_4bit_stop_sample.py

# 生成トークンの停止条件設定クラス
class MyStoppingCriteria(StoppingCriteria):
    def __init__(self, stop_words_ids=[]):
      super().__init__()
      self.stop_words_ids = stop_words_ids
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
        for stop_word_ids in self.stop_words_ids: 
            # 生成した末尾のトークンがストップワードと一致するか確認
            if input_ids[0][-1] == stop_word_ids: 
                # トークンの生成をストップ
                return True
        return False

# ストップワードのリスト
stop_words = [" machine", " human"]

stop_words_ids = []
for stop_word in stop_words:
    # 文字列をテンソルに変換
    stop_word_ids = tokenizer(stop_word, return_tensors='pt', add_special_tokens=False)['input_ids']
    # テンソルをcudaに配置
    stop_word_ids = stop_word_ids.to("cuda:0") 
    # テンソルをリストに格納
    stop_words_ids.append(stop_word_ids) 

# MyStoppingCriteriaのインスタンスを生成
my_stopping_criteria = MyStoppingCriteria(stop_words_ids=stop_words_ids) 

# StoppingCriteriaListのインスタンスを生成
stopping_criteria_list = StoppingCriteriaList([my_stopping_criteria]) 

# パイプラインの設定
pipe = pipeline(
    "text-generation", # タスクの指定
    model=model, # モデルをセット
    tokenizer=tokenizer, # トークナイザーをセット
    max_new_tokens=64, # 生成するトークンの最大数
    stopping_criteria=stopping_criteria_list # StoppingCriteriaListのインスタンスをセット
)

上記の流れを言葉にすると以下のようになります。
StoppingCriteriaクラスを継承したMyStoppingCriteriaクラスを定義します。
↓
ストップワード（今回は" machine"と" human"を設定）をテンソルに変換。
↓
MyStoppingCriteriaクラスに（テンソルに変換した）ストップワードを渡してインスタンス化。
↓
StoppingCriteriaListにMyStoppingCriteriaのインスタンスを渡してインスタンス化。
↓
pipelineにStoppingCriteriaListのインスタンスを渡す。

以上がストップワード設定の手順となります。

主に停止する際の条件はMyStoppingCriteriaクラス内に書きます。__call__メソッドがトークン（単語や記号）を生成するたびに呼び出され、Trueを返すとトークンの生成が停止します（今回の例だと" human" or " machine"のトークンが生成されると停止します）。

出力確認

Artificial intelligence (AI) is the ability of a machine

pipe_4bit_stop_sample.py

stop_words = [" machine", " human"]

上記、" machine"、" human"の頭にある空白「" "」を削除すると停止できません。前の単語との間にある空白を考慮する必要があります。
tokenizerや場合にもよる（"Artificial"や"AI"はOK）ので、以下のように空白の有り無しなど複数のバリエーションを設定しておくと良いかもしれません。

pipe_4bit_stop_sample.py

stop_words = ["machine", " machine", "human", " human"]

おまけ

文章（複数ワード）の一致に対応したストップワードの設定

MyStoppingCriteriaを以下のように変更することで文章（複数ワード）の一致に対応できます。

pipe_4bit_stop_multiple_sample.py

# 生成トークンの停止条件設定クラス
class MyStoppingCriteria(StoppingCriteria):
    def __init__(self, stop_words_ids=[]):
      super().__init__()
      self.stop_words_ids = stop_words_ids
      # テンソルの要素数のリストを作成
      self.stop_words_ids_count = [stop_word_ids.numel() for stop_word_ids in self.stop_words_ids]
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
        for stop_word_ids, stop_word_ids_count in zip(self.stop_words_ids, self.stop_words_ids_count):
            # 末尾の複数トークン（ストップワードと同じ要素数）がストップワードと一致するか確認
            if torch.all(input_ids[0][-stop_word_ids_count:] == stop_word_ids):
                return True
        return False

全体のコード例は以下になります。

pipe_4bit_stop_multiple_sample.py

import os
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
    BitsAndBytesConfig,
    StoppingCriteria,
    StoppingCriteriaList
)
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_core.prompts import PromptTemplate
from langchain.chains.llm import LLMChain

# huggingfaceのトークンの設定（gemmaを使用するのに必要なため）
os.environ["HF_TOKEN"] = '<< huggingfaceのトークン >>'

# gemma-2b-it のリポジトリIDを設定
repo_id = "google/gemma-2b-it"

# 量子化の構成を設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True, # 4bit量子化を使用
    bnb_4bit_quant_type="nf4", # 量子化データ型にNF4を設定
    bnb_4bit_compute_dtype=torch.bfloat16 # 量子化計算時のデータ型にbfloat16を設定
)

# モデルの設定
model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    device_map={"": 'cuda:0'}, # gpu（cuda）での実行を指定
    quantization_config=quantization_config # 量子化の構成をセット
)

# トークナイザーの設定
tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path=repo_id,
    use_fast = False,
)

# 生成トークンの停止条件設定クラス
class MyStoppingCriteria(StoppingCriteria):
    def __init__(self, stop_words_ids=[]):
      super().__init__()
      self.stop_words_ids = stop_words_ids
      # テンソルの要素数のリストを作成
      self.stop_words_ids_count = [stop_word_ids.numel() for stop_word_ids in self.stop_words_ids]
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
        for stop_word_ids, stop_word_ids_count in zip(self.stop_words_ids, self.stop_words_ids_count):
            # 末尾の複数トークン（ストップワードと同じ要素数）がストップワードと一致するか確認
            if torch.all(input_ids[0][-stop_word_ids_count:] == stop_word_ids):
                return True
        return False

# ストップワードのリスト
stop_words = [" human intelligence", " problem-solving"] 

stop_words_ids = []
for stop_word in stop_words:
    # 文字列をテンソルに変換
    stop_word_ids = tokenizer(stop_word, return_tensors='pt', add_special_tokens=False)['input_ids']
    # テンソルをcudaに配置
    stop_word_ids = stop_word_ids.to("cuda:0") 
    # テンソルをリストに格納
    stop_words_ids.append(stop_word_ids) 

# MyStoppingCriteriaのインスタンスを生成
my_stopping_criteria = MyStoppingCriteria(stop_words_ids=stop_words_ids) 

# StoppingCriteriaListのインスタンスを生成
stopping_criteria_list = StoppingCriteriaList([my_stopping_criteria]) 

# パイプラインの設定
pipe = pipeline(
    "text-generation", # タスクの指定
    model=model, # モデルをセット
    tokenizer=tokenizer, # トークナイザーをセット
    max_new_tokens=64, # 生成するトークンの最大数
    stopping_criteria=stopping_criteria_list # StoppingCriteriaListのインスタンスをセット
)

# langchainで扱える形式に変換
llm = HuggingFacePipeline(
    pipeline=pipe # パイプラインをセット
)

# プロンプトテンプレートをフォーマット
question_prompt_template_format = tokenizer.apply_chat_template(
    conversation = [
        {"role": "user", "content": "{question}"}
    ], 
    tokenize=False, 
    add_generation_prompt=True
)

# プロンプトの設定
QUESTION_PROMPT = PromptTemplate(
    template=question_prompt_template_format,
    input_variables=["question"]
)

# 質問回答chainの設定
chain = LLMChain(llm=llm, prompt=QUESTION_PROMPT)

# 質問文（和訳: AIについて50字程度で簡潔に説明してください。）
question = 'Briefly describe AI in 50 words or so.'

# LLMの回答生成
response = chain.invoke({"question": question})

# 回答を確認
print(response["text"])

出力確認

Artificial intelligence (AI) is the ability of a machine to perform tasks that typically require human intelligence

おわりに

モデルをlangchainに繋げられればRAGなどが活用しやすくなります。また、量子化は速度の向上、ストップワードの設定は同じ単語の繰り返しの防止などを期待できます。

次もllm関連の記事になる予定です。また、機会があればよろしくお願いします。

huggingfaceのモデルをlangchainで使う＋量子化とストップワード（停止文字列）の設定

はじめに

gguf形式のモデルを扱う場合

HuggingFacePipelineを使った実行

コーデイング

出力確認

HuggingFacePipelineを使った実行（量子化）

コーディング

出力確認

HuggingFacePipelineを使った実行（量子化 + ストップワード）

コーディング

出力確認

おまけ

文章（複数ワード）の一致に対応したストップワードの設定

出力確認

おわりに

参考

Discussion