RWKVでembedding vectorを計算

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)

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])

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()