RWKVでembedding vectorを計算

Kan Hatakeyama

概要

背景

  • 最近は､GPTが流行ってます

  • しかしGPT-3.5以降はfine tuningが執筆時点でできません

  • なので､オリジナルデータを学習させるには､少し工夫が必要です

  • 要するに､文章のembedding vectorを計算する必要があります

  • しかし､GPTのAPIは地味に値段が高いため､pdfが100個くらいあったりすると､破産する恐れが出てきます

目的

  • 最終的な推論はGPT-3.5 or 4にやらせるとして､embedding vectorの計算は､もう少しローコスト･低性能なLLMで良いのではないかと､誰しも考えるはずです

    • 或いは､google検索のような､index検索を使うのも手です

      • ただしこの場合は､言語を跨いだ検索などが難しそうです

  • そこで､RNNでTransformer並みの動きをすると話題のRWKVで､embedding vector的なものを計算してみようと考えました

結果

こちらでRWKVを入れておきます｡とりあえず7Bモデルを使ってみました｡14Bモデルの方が高性能なはずです｡

モデル読み込み

import os
# set these before import RWKV
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0' # '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries

from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
 #model読み込み 
model = RWKV(model='モデルパス', strategy='cpu fp32')
pipeline = PIPELINE(model, "20B_tokenizer.json") # 20B_tokenizer.json is in https://github.com/BlinkDL/ChatRWKV

デモコードの動作確認

 #与えられたtextから ､その続きを生成する

ctx = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."
ctx="How are you?"
ctx="こんにちは"
print(ctx, end='')

def my_print(s):
    print(s, end='', flush=True)


args = PIPELINE_ARGS(temperature = 1.0, top_p = 0.7, top_k = 100, # top_k = 0 then ignore
                     alpha_frequency = 0.25,
                     alpha_presence = 0.25,
                     token_ban = [0], # ban the generation of some tokens
                     token_stop = [], # stop generation whenever you see any token here
                     chunk_len = 256) # split input into chunks to save VRAM (shorter -> slower)

pipeline.generate(ctx, token_count=200, args=args, callback=my_print)

結果はランダムですが､怪しい文字列が出てきました｡五年ほどにRNNを走らせた際､滅茶苦茶な文字列が出てきた記憶がよみがえりました｡

こんにちは！ 今回の記事はバングラデシュからの入国を目的としたアクセスを無制限にしてきた人物の「逮捕・拘留・不当送還」ということです。今回はこの件について触れます。

Vector計算
model.forwardにencode済みのtextを投げると､出力(out)とモデルの状態(state)が返って来ます｡GPTと違い､ca. 1500次元のベクトルが返ってくることはなく､160 x 4096 = 655360 次元の高次元なベクトルがそのまま返って来ます｡ ベクトル検索ライブラリに組み込むのであれば､何らかの圧縮作業が必要かもしれません

import numpy as np
def calc_vec(txt,model):
    inp_vec=pipeline.encode(txt)
    out, state = model.forward(inp_vec, None)
    np_state=[i.detach().cpu().numpy() for i in state]
    np_state=np.concatenate(np_state, axis=0)
    return np_state

実際に計算させてみます｡適当に思い付いた文字列を渡します｡

text_list=[
"吾輩は猫である",
"汝の名は女だ",
"I am a cat",
"ピペットで試薬を三回分取した",
"アセトンは有機溶媒である",
"非プロトン性溶媒として､THF､トルエンなどが挙げられる",
"トルエンをこぼしてしまいました",
"私は猫です",
"私は犬です",
"日本的霊性は，鎌倉時代に禅と浄土系思想によって初めて明白に顕現し，その霊性的自覚が現在に及ぶと述べる．",
]

from tqdm import tqdm
vec_list=[calc_vec(i,model) for i in tqdm(text_list)]

target_vec=vec_list[0]

from scipy import spatial
def cos_sim(a, b):
    return 1 - spatial.distance.cosine(a, b)


cos_sim_list=[cos_sim(target_vec,i) for i in vec_list]

for i in np.argsort(cos_sim_list)[::-1]:

    print("{:.3f}".format(cos_sim_list[i]),text_list[i])

｢吾輩は猫である｣とコサイン類似度の高い文字列は以下の通り

1.000 吾輩は猫である
0.950 汝の名は女だ
0.950 アセトンは有機溶媒である
0.946 トルエンをこぼしてしまいました
0.942 私は猫です
0.941 ピペットで試薬を三回分取した
0.917 私は犬です
0.905 非プロトン性溶媒として､THF､トルエンなどが挙げられる
0.882 I am a cat
0.859 日本的霊性は，鎌倉時代に禅と浄土系思想によって初めて明白に顕現し，その霊性的自覚が現在に及ぶと述べる．

残念ながら､｢私は猫です｣､｢I am a cat｣は上の方に来ませんでした｡
コンテンツというよりは､文語調 or 口語調の違いをLLMが察してしまった気がします｡

次に､ I am a catで計算してみます｡

1.000 I am a cat
0.914 私は猫です
0.911 私は犬です
0.899 汝の名は女だ
0.882 吾輩は猫である
0.849 トルエンをこぼしてしまいました
0.841 アセトンは有機溶媒である
0.826 ピペットで試薬を三回分取した
0.768 非プロトン性溶媒として､THF､トルエンなどが挙げられる
0.706 日本的霊性は，鎌倉時代に禅と浄土系思想によって初めて明白に顕現し，その霊性的自覚が現在に及ぶと述べる．

無事に､｢私は猫です｣が､類似の高い結果として返って来ました｡
素晴らしい!
100%の精度ではありませんが､まあ使いようのあるLLMだと思います｡

2023/4/27 追記

軽量モデルRWKV-4b-Pile-3Bでも計算してみました｡あまり精度は変わらないですね｡

1.000 I am a cat
0.828 私は犬です
0.813 私は猫です
0.787 汝の名は女だ
0.760 吾輩は猫である
0.690 トルエンをこぼしてしまいました
0.688 アセトンは有機溶媒である
0.642 ピペットで試薬を三回分取した
0.560 非プロトン性溶媒として､THF､トルエンなどが挙げられる
0.447 日本的霊性は，鎌倉時代に禅と浄土系思想によって初めて明白に顕現し，その霊性的自覚が現在に及ぶと述べる．

MDSで次元削減

高次元ベクトルを2次元に転写してみます

from sklearn.manifold import MDS
reducer=MDS(n_components=2)
comp_vec=reducer.fit_transform(np.array(vec_list))

import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'MS Gothic'
plt.scatter(comp_vec[:,0],comp_vec[:,1])
for i, txt in enumerate(text_list):
    plt.annotate(txt, (comp_vec[i,0],comp_vec[i,1]),size=10)
    
plt.show()

そこそこの精度でしょうか｡

まとめ

RWKVでembedding vector的なものを計算してみました｡悪くない精度だと思います｡本手法に加え､index検索も併用すれば､欲しい情報は概ねfetch出来るように思いました｡vicunaよりは軽量なので､このようなタスクでは使い道があるかもしれません｡

この記事が気に入ったらサポートをしてみませんか？