AnyTech Engineer Blog

AnyTech Engineer Blogは、AnyTechのエンジニアたちによる調査や成果、Tipsなどを公開するブログです。

JARVIS(っぽい何か)を作ろう!第二回:音声認識した結果から返答生成する

JARVIS(っぽい何か)を作ろう!第二回:音声認識した結果から返答生成する

はじめに

本記事はおしゃべりAIをオフラインかつローカルで実装するシリーズ第二回です。 今回は、text to textという手法を用いて与えられた質問(テキスト)から返事(テキスト)を生成するプログラムを書きます。

前回:JARVIS(っぽい何か)を作ろう!第一回:マイクからリアルタイムで音声認識する

シリーズ

こちらの手順を踏まえて、作っていきたいと思います。 ※ 本シリーズはローカルUbuntuマシンにGPUがある場合を想定しております。

第一回:マイクからリアルタイムで音声認識する 第二回:音声認識した結果から返答生成する ← 今回はこちらになります。 第三回:返答を合成音声で喋らせる 第四回:キャラクターが喋っているような見た目を作る 第五回:今までで作ったものを組み合わせる

今回の目次

  1. text to textモデルについて
  2. text to textモデルの実装
  3. 前回の音声認識の結果をtext to textのモデルに渡す

1. text to textモデルについて

概要

text to textモデルは「自然言語処理モデル」や「Text2Text」、最近ではChatGPTの登場で特にパラメータの多いtext to textモデルは「LLM(Large Language Models)」とも呼ばれています。 これらモデルは大量のテキストデータを学習した自然言語処理モデルのことで、主に次の用途で用いられます。

  • テキスト生成
    • 質疑応答
    • 文章要約
    • 翻訳
  • テキスト分類
    • 言語分類
    • 感情分類

...などがあります。 同じモデルでも、学習方法によって用途にあった出力を得られるように使用されます。

モデルの例

現在、我々が使用できるモデルはAPI経由で使用できるもの(GPT3.5Turboなど)と、ローカル環境で動作するものがあります。 今回は自宅のPCでなんとか動く範囲の、日本語が学習データに含まれているモデルをいくつかご紹介します。

rinna/japanese-gpt2-medium

RTX3090があれば、推論と学習が十分に可能。

rinna/japanese-gpt-neox-3.6b

RTX3090で推論が可能(他のモデルも実装するならfp16が必須か)。筆者は学習未実施。

cyberagent/open-calm-3b

RTX3090で推論が可能(他のモデルも実装するならfp16が必須か)。筆者は学習未実施。

BlinkDL/rwkv-4-raven

RTX3090で推論が可能(VRAM24以下ならfp16必須)。筆者は学習未実施。 オプションをいくつか使用すれば、かなりの推論速度が出る。

今回はRWKV(ルワクフと読む模様です)をメインに使用していきます。

2. text to textモデルの実装

環境構築

RWKVの一部オプションはCUDAのバージョンが指定以上でないと動かないので、こちらはDocker上に構築します。

Dockerfile

# RTX3090
FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel

ENV DEBIAN_FRONTEND noninteractive
ARG project_dir=/app/
WORKDIR $project_dir

ENV PYTHON_VERSION 3.8.1
ENV HOME /root
ENV PYTHON_ROOT $HOME/python-$PYTHON_VERSION
ENV PATH $PYTHON_ROOT/bin:$PATH
ENV PYENV_ROOT $HOME/.pyenv

RUN apt-get update && apt-get upgrade -y \
 && apt-get install -y \
    git \
    make \
    build-essential \
    libssl-dev \
    zlib1g-dev \
    libbz2-dev \
    libreadline-dev \
    libsqlite3-dev \
    wget \
    curl \
    llvm \
    libncurses5-dev \
    libncursesw5-dev \
    xz-utils \
    tk-dev \
    libffi-dev \
    liblzma-dev

RUN git clone https://github.com/pyenv/pyenv.git $PYENV_ROOT
RUN $PYENV_ROOT/plugins/python-build/install.sh
RUN /usr/local/bin/python-build -v $PYTHON_VERSION $PYTHON_ROOT
RUN rm -rf $PYENV_ROOT

RUN pip install --upgrade pip
RUN pip install numpy
RUN pip install transformers
RUN pip install datasets
RUN pip install sentencepiece
RUN pip install -U protobuf
RUN pip install dill
RUN pip install langchain==0.0.61

Dockerfileをビルドして、runします。

docker build . -t text2text
docker run -v $PWD:/app \
       --privileged \
       --device /dev/video0:/dev/video0:mwr \
       --device /dev/snd:/dev/snd \
       --device /dev/usb:/dev/usb \
       -v /run/dbus/:/run/dbus/:rw \
       -v /dev/shm:/dev/shm \
       --shm-size=10g --runtime=nvidia --name text2text_run --rm -it text2text bash

推論コードの実装

EngAndMoreには日本語データも入っているので、こちらを使用します。

wget https://huggingface.co/BlinkDL/rwkv-4-raven/blob/main/RWKV-4-Raven-3B-v8-EngAndMore-20230408-ctx4096.pth

ダウンロードしたモデルファイルをmodels/rwkvに保存します。

tools./text_generation.py

import os
import sys
os.path.join(os.getcwd())
sys.path.append('./')
from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
from utils.prompt_utils import generate_prompt
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '1'
model_path = 'models/rwkv/RWKV-4-Raven-3B-v8-EngAndMore-20230408-ctx4096.pth'

# プロンプト
input=""

class RWKV4Raven():
    def __init__(self):
        self.input = input
        self.last_input = ""
        self.model = RWKV(model=model_path, strategy='cuda fp16')
        self.pipeline = PIPELINE(self.model, "models/rwkv/20B_tokenizer.json")
        self.args = PIPELINE_ARGS(temperature = 1.0, top_p = 0.999, top_k = 100,
                                  alpha_frequency = 0.25, alpha_presence = 0.25,
                                  token_ban = [0],
                                  token_stop = [],
                                  chunk_len = 512)

    def generate(self, msg):
        res = ""
        if len(msg) > 0:
            res = self.pipeline.generate(generate_prompt(instruction=msg, input=self.input),
                                         token_count=80, args=self.args)
            # お好み後処理、改行前の文のみ
            if "\n" in res:
                res = res.split("\n")[0]
        return res

rwkv = RWKV4Raven()
msg = "おはようございます"
res = rwkv.generate(msg)
print(f'res = {res}')

実行してみます。

python tools/text_generation.py

エラーが出てしまいました。

AttributeError: 'RecursiveScriptModule' object has no attribute 'cuda_att_seq'

RWKVのソースに追記

RWKVのソース元を取ってきて試してみます。

cd models
git clone https://github.com/BlinkDL/ChatRWKV.git
cd ../

クローンしてきたソースmodels/ChatRWKV/rwkv_pip_package/src/rwkv/model.pyの12行目に、以下を追加します。

os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '1'

また、python tools/text_generation.pyの一部を変更します。CloneしてきたRWKVのソースを使うように変更します。

#from rwkv.model import RWKV
from models.ChatRWKV.rwkv_pip_package.src.rwkv.model import RWKV

もう一度python tools/text_generation.pyを実行

再度python tools/text_generation.pyを実行すると、以下のエラーが出ます。

Traceback (most recent call last):
  File "tools/text_generation.py", line 6, in <module>
    from models.ChatRWKV.rwkv_pip_package.src.rwkv.model import RWKV
  File "./models/ChatRWKV/rwkv_pip_package/src/rwkv/model.py", line 30, in <module>
    load(
  File "/root/python-3.8.1/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1284, in load
    return _jit_compile(
  File "/root/python-3.8.1/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1509, in _jit_compile
    _write_ninja_file_and_build_library(
  File "/root/python-3.8.1/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1593, in _write_ninja_file_and_build_library
    verify_ninja_availability()
  File "/root/python-3.8.1/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1649, in verify_ninja_availability
    raise RuntimeError("Ninja is required to load C++ extensions")
RuntimeError: Ninja is required to load C++ extensions

必要なライブラリをインストール

Ninjaをインストールします。

wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
mkdir pkg
mv ninja-linux.zip pkg/
apt install unzip
unzip pkg/ninja-linux.zip -d /usr/local/bin/
update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force

もう一度python tools/text_generation.pyを実行

再度、python tools/text_generation.pyを実行します。 RWKVモデルから返答が返ってきました。

res = おはようございます。どうぞご安心いただくようお願いします。何かお伺いしたいことがあれば教えてください。お時間を限られている時間に助けてください。

3. 前回の音声認識の結果をtext to textのモデルに渡す

USBマイクとDocker環境が相性が良くないので、別々のプログラムに分けて組み合わせることにしました。 こんな流れで実装します。 | speech recognition (ローカル環境)| ↓ UDP通信で音声認識結果のテキストを送信 ↓ | text to text (Docker環境)|

フォルダ構成

├── Dockerfile
├── models
│   ├── ChatRWKV
│   ├── rwkv
│   │   ├── 20B_tokenizer.json
│   │   ├── RWKV-4-Raven-3B-v8-EngAndMore-20230408-ctx4096.pth
├── pkg
│   └── ninja-linux.zip
├── scripts
│   └── docker_run.sh
├── tools
│   ├── text_generation.py
│   └── voice_recognize.py
└── utils
    └── prompt_utils.py

準備

まず、Dockerとの通信用にネットワークを用意します。ローカル環境で以下を実行します。

docker network create --subnet=192.168.66.0/24 jarvis_app

docker network lsで、以下があればOKです。

NETWORK ID     NAME                   DRIVER    SCOPE
819djad1c72a1   jarvis_app             bridge    local

音声認識の実行と実装(前回からUDPを追加しています)

ローカル環境で、python tools/voice_recognition.pyを実行します。

python tools/voice_recognition.py

import os
import sys
import soundfile
import pyaudio
import wave
import numpy as np
import socket
from contextlib import closing
from espnet_model_zoo.downloader import ModelDownloader
from espnet2.bin.asr_inference import Speech2Text
#apt install portaudio19-dev
FORMAT = pyaudio.paInt16

# Audio params
chunk = 1024  #1つの音声チャンクの長さを1024サンプルとする
threshold = 0.50  #マイクの音量がこれを超えると録音を開始する
sampling_rate = 16000  #サンプリングレート、マイク性能に依存
record_seconds = 15  #録音時間の上限
record_cache_seconds = 0.9
stop_seconds = 2  #録音中断までの時間
B = int(sampling_rate / chunk * stop_seconds)
A = int(sampling_rate) / int(chunk) * int(record_seconds)
A_CACHE = int(sampling_rate) / int(chunk) * int(record_cache_seconds)

machine_list = ['docker']
ip_list = ['192.168.66.66']
module_list = ['t2t']
port_list = [13103]


def send_text(text, host, port):
    sock = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
    with closing(sock):
        message = '{}'.format(text).encode('utf-8')
        sock.sendto(message,(host,port))
    return


def load_voice_recog_model():
    d = ModelDownloader()
    model = Speech2Text(
        **d.download_and_unpack("Shinji Watanabe/laborotv_asr_train_asr_conformer2_latest33_raw_char_sp_valid.acc.ave"),
        device="cuda"
    )
    return model


def save_wave(all, audio_path):
    data = b''.join(all)
    out = wave.open(audio_path,'w')
    out.setnchannels(1) #mono
    out.setsampwidth(2) #16bits
    out.setframerate(sampling_rate)
    out.writeframes(data)
    out.close()


def inference(speech2text, audio_path):
    speech, _ = soundfile.read(audio_path)
    nbests = speech2text(speech)
    text, *_ = nbests[0]
    return text


def recognize(speech2text):
    audio_path = "speech.wav"
    p = pyaudio.PyAudio()
    stream = p.open(format=FORMAT, channels=1, rate=sampling_rate, input=True, frames_per_buffer=chunk)
    all = []
    while True:
        #まずサンプルを取る
        data = stream.read(chunk, exception_on_overflow=False)
        x = np.frombuffer(data, dtype="int16") / 32768.0
        xmax = x.max()
        print(xmax)
        if xmax <= threshold:
            all = []
            for i in range(0, int(A_CACHE)):
                data_cache = stream.read(chunk, exception_on_overflow=False)
                all.append(data_cache)
        else:
            n = 0  #連続して閾値を下回った回数を記録するための変数
            for i in range(0, int(A)):
                data = stream.read(chunk, exception_on_overflow=False)
                all.append(data)
                x = np.frombuffer(data, dtype="int16") / 32768.0
                xmax = x.max()
                if xmax <= threshold:
                    n += 1  #閾値を下回っていたらincrementする
                else:
                    n = 0  #閾値を上回ったらリセットする
                # 中断判定
                if n == B:
                    break
            # 音声保存
            save_wave(all, audio_path)
            all = []

            # 音声認識
            text = inference(sr_model, audio_path)
            print(f"Recognized text = {text}")

            # 認識した文字数が3以上ならtext to textモデル環境に送信する
            if len(text) > 2:
                send_text(text,
                      ip_list[machine_list.index('docker')],
                      port_list[module_list.index('t2t')])

if __name__ == '__main__':
    sr_model = load_voice_recog_model()
    recognize(sr_model)

docker環境準備

docker環境から一度出て、192.168.66.66でdockerを立ち上げます。

docker run -v $PWD:/app \
       --privileged \
       --device /dev/video0:/dev/video0:mwr \
       --device /dev/snd:/dev/snd \
       --device /dev/usb:/dev/usb \
       -v /run/dbus/:/run/dbus/:rw \
       -v /dev/shm:/dev/shm \
       --net=jarvis_app --ip=192.168.66.66 \
       --shm-size=10g --runtime=nvidia --name text2text_run --rm -it text2text bash

再度、Ninjaをインストールします。

apt install unzip
unzip pkg/ninja-linux.zip -d /usr/local/bin/
update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force

text to textの実行と実装(UDP受信を実装)

Docker環境で、python tools./text_generation.pyを実行します。

tools./text_generation.py

import os
import sys
from socket import socket, AF_INET, SOCK_DGRAM
os.path.join(os.getcwd())
sys.path.append('./')
from models.ChatRWKV.rwkv_pip_package.src.rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
from utils.prompt_utils import generate_prompt
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '1'
model_path = 'models/rwkv/RWKV-4-Raven-3B-v8-EngAndMore-20230408-ctx4096.pth'

module_list = ['t2t']
port_list = [13103]

# プロンプト
input=""

class RWKV4Raven():
    def __init__(self):
        self.input = input
        self.last_input = ""
        self.model = RWKV(model=model_path, strategy='cuda fp16')
        self.pipeline = PIPELINE(self.model, "models/rwkv/20B_tokenizer.json")
        self.args = PIPELINE_ARGS(temperature = 1.0, top_p = 0.999, top_k = 100,
                                  alpha_frequency = 0.25, alpha_presence = 0.25,
                                  token_ban = [0],
                                  token_stop = [],
                                  chunk_len = 512)

    def generate(self, msg):
        res = ""
        if len(msg) > 0:
            res = self.pipeline.generate(generate_prompt(instruction=msg, input=self.input),
                                         token_count=80, args=self.args)
            # お好み後処理、改行前の文のみ
            if "\n" in res:
                res = res.split("\n")[0]
        return res


if __name__ == '__main__':
    # Set port setting to ear module
    s = socket(AF_INET, SOCK_DGRAM)
    s.bind(('', port_list[module_list.index('t2t')]))

    # Load model
    rwkv = RWKV4Raven()

    print('=> Waiting for speech recognition...')
    while True:
    # Receive msg
    msg, address = s.recvfrom(8192)
    msg = msg.decode("utf-8")

        # Inference
        res = rwkv.generate(msg)
    print(f'msg = {msg}')
    print(f'res = {res}')

結果

いい感じにできました!

=> Waiting for speech recognition...
msg = おはようございます
res = おはようございます。どうかな?
msg = 今日の天気はどうですか
res = 今日の天気には、集中力を保つために一日中体力を働かなければならない場合があります。そのため、筋肉の質をより優れる方法を探すために、毎日行軍するのがおすすめです。