概要

HuggingFace Hubに登録されているモデルをローカルにダウンロードして、LangChain経由で対話型のプログラムを作成する。

前提条件

ランタイムは Python 3.11.6 を使用する。今回使用するモデルには、比較的汎用的な計算機でも動作する "TinyLlama/TinyLlama-1.1B-Chat-v1.0" を採用した

使用する Python パッケージは以下の通り。

パッケージ名	バージョン
langchain	0.0.322
transformers	4.35.2
huggingface_hub	0.20.2

実装コード

実装は以下の通り。

snapshot_download関数でリポジトリのスナップショットを撮り、そのスナップショットにあるファイルを元に tokenizer と model を準備する。次にHuggingFacePipelineクラスをインスタンス化し、これとChatPromptTemplateでチェーンを作成し、chain.invokeメソッドで実行する。

app.py

from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from langchain.prompts import ChatPromptTemplate
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from huggingface_hub import snapshot_download

# 1.HuggingFaceからモデルをダウンロード
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
download_path = snapshot_download(repo_id=model_id)

# 2.TokenizerとModelをスナップショットからインスタンス化
tokenizer = AutoTokenizer.from_pretrained(download_path)
model = AutoModelForCausalLM.from_pretrained(download_path)

# 3.HuggingFacePipelineを作成
# ref. https://python.langchain.com/docs/integrations/llms/huggingface_pipelines
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
hf = HuggingFacePipeline(pipeline=pipe)

# 4.ChatPromptTemplateを作成
prompt = ChatPromptTemplate.from_template(
    "<|system|>You are a chatbot who can help code!</s><|user|>{question}</s><|assistant|>"
)
chain = prompt | hf

# 5.Chainの実行
question = "Write me a function to print out 'Hello, World!' to the CLI in Python."
answer = chain.invoke({"question": question})
print(answer)

実行結果

実行した結果は以下の通り。
スナップショットから "TinyLlama/TinyLlama-1.1B-Chat-v1.0" のデータをフェッチし、LangChainを介してモデルとの対話を行えている事がわかる。

bash

$ python app.py
Fetching 10 files: 100%|█████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 10860.45it/s]

Certainly! Here's a Python function that prints 'Hello, World!' to the CLI:

```python
def print_hello_world():
    print("Hello, World!")
```

To use this function, simply call it with no arguments:

```python
print_hello_world()
```

This will print 'Hello, World!' to the CLI. Note that this function assumes that you have already imported the `print` function from the `sys` module. If you're not using Python 3.6 or later, you may need to import the `print` function explicitly:

```python
from sys import print as print_
```

補足：GPUによる推論時間の短縮

GPUが利用できるのであれば、引数にdevice=0を指定することで、計算時間をかなり短縮できる。

app.py (抜粋)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256, device=0)

以下は i7-13700F + RTX 4070 の構成で、10回推論を行ったときの計算時間の平均である。
※ かなり雑な計測なので、実際に手元で試してほしい

計測ケース	計測時間（秒）
CPU のみ	24.3
GPU サポート	2.3

参考文献