📈

Hugging Face Spacesで、簡単ChatBotをつくってみる

2024/05/16に公開

Streamlit

目的

こちらの記事の続きで、もう少し違ったStreamlitアプリを構築したいと思います。
前回は、画像解析だったので、今回はChatBotにします。

構築するもの

HuggingChatという素敵な見た目（ChatGPTそのもの）な、Streamlitサービスがあります。
https://huggingface.co/chat/

そのコードを参考にしつつ、裏側には自由にモデルを指定したり、目的を絞ったChatBotにしたいと思います。

分かり易いものとして、ChatBotで聞いたら今日の株価を教えてくれるようなものを考えます。

手順

環境構築や、HuggingFaceSpacesでの操作は前回記事を参照ください。

requirements.txtだけ一部異なるので、下記をご利用ください。

streamlit
sentencepiece
transformers
bs4
Cython
torch

プログラム作成

少し長いプログラムになったので、パーツに分けて解説していきます。

まずは、Transformersのモデル・パイプラインを指定しています。
今回は、ChatBotなので、QuestionAnsweringのパイプラインを使います。

モデルは、軽くてある程度の性能が見込めている'KoichiYasuoka/bert-base-japanese-wikipedia-ud-head'のモデルを拝借します。

from transformers import AutoTokenizer
from transformers import AutoModelForQuestionAnswering
from transformers import QuestionAnsweringPipeline

model_name = 'KoichiYasuoka/bert-base-japanese-wikipedia-ud-head'

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
qap = QuestionAnsweringPipeline(tokenizer=tokenizer, model=model)

続いて、株価の情報を取ってくるところです。
ここでは、スクレイピングで情報を取ってきましょう。

まずは、streamlitのtext_inputでユーザーが株コードを指定するように作っています。
それが指定されて、buttonを押下すると、スクレイピングが実行されてcontentの中に情報が格納されます。

ただし、contentのままだと、streamlit上で別の動作をすると情報がなくなってしまうので、st.session_state.contentという風に、streamlitのセッション情報に残しておきます。

import re
import requests
import unicodedata
from bs4 import BeautifulSoup

import streamlit as st

st.title("株価お知らせBot")
stock_code = st.text_input('株コード', placeholder='株コード', max_chars=4, help='4桁の数字')

if "content" not in st.session_state: 
    st.write("はじめに、株コードを指定してください。")

if st.button("株決定"):
    url = f"https://www.nikkei.com/nkd/company/?scode={stock_code}"
    res = requests.get(url)
    soup = BeautifulSoup(res.text, 'html.parser')

    _text = soup.find('div', attrs={ 'class': 'm-stockInfo_top_left' })
    content = _text.text

    content = unicodedata.normalize('NFKD', content)
    st.session_state.content = re.sub('[\r\t\n]+', ' ', content)

続いて、上記で指定したパイプラインと、スクレイピングした情報から、回答を生成する部分です。
変数answerは、dict型で「{'score': 0.9943368434906006, 'start': 11, 'end': 18, 'answer': '回答はこちら'}」のような返り値になるため、answer部分のみを抽出しています。

# Transformersで回答を作成
def generate_response(prompt, max_length=50):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")

    # Generate response
    answer = qap(context=st.session_state.content[:100], question=prompt)
    
    return answer["answer"]

最後は、ChatBotをStreamlit上で実装している部分です。
基本は、上述の通りHugging Chatのコードを真似させていただいています。

ここまで簡単なコードでChatGPT風のChatBotの枠が作れてしまうのは、Streamlit恐るべしです。

# メッセージがない時
if "messages" not in st.session_state.keys():
    st.session_state.messages = [{"role": "assistant", "content": "何か御用ですか？"}]

# チャット内容の表示
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# ユーザーの質問
if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

# AIによる回答
if st.session_state.messages[-1]["role"] != "assistant":
    with st.chat_message("assistant"):
        with st.spinner("考え中..."):
            response = generate_response(prompt)
            st.write(response) 
    message = {"role": "assistant", "content": response}
    st.session_state.messages.append(message)

実行

HuggingFaceSpacesにPushする方法は、こちらを見ていただくとして、最後に出力イメージを確認しておきましょう。

まずは初期画面です。

株コードを入れて、「株決定」を押すと、スクレイピングが動きます。
今回は、NTTの株コードを入れてみています。

少し待つと、RUNNINGが止まるので、さっそく質問してみましょう。

高値や、始値について教えてくれています。
ただ、「157.6 円安値」のように、スクレイピング元のサイトの情報での切れ目をモデルが理解できずに回答してしまっていますね。

この辺りは、contextとして渡す情報を、もう少し見直す必要があるのか。
または、別のモデルやパイプラインを使えば解決するのか、今後検討していきたいと思います。

目的

構築するもの

手順

プログラム作成

実行

Discussion