rwkv-emb 0.0.5

The EmbeddingRWKV Model

The EmbeddingRWKV Model

https://github.com/howard-hou/EmbeddingRWKV

Tokenizer

import os
from rwkv_emb.tokenizer import RWKVTokenizer

# 1. 
tokenizer = RWKVTokenizer()

# 2. add_eos should be True for embedding inference to append the EOS token (65535)
text = "hello world"
tokens = tokenizer.encode(text, add_eos=True)

print(tokens)
# 输出: [token_id, ..., 65535]

The encode method returns a list of integers. For embedding inference, append the end-of-sequence token (65535) to mark completion before feeding the tokens to the model.

# !!! set these before import RWKV !!!
import os

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

from rwkv_emb.model import EmbeddingRWKV

EOS_INDEX = 65535

# download models: to be announced
model = EmbeddingRWKV(model_path='path-to-model')

# !!! model.forward(tokens, state) will modify state in-place !!!
# single-sample inference
emb, state = model.forward([187, 510, 1563, 310, 247, EOS_INDEX], None)
print(emb.detach().cpu().numpy())                   # get logits

# streaming a single sequence
emb, state = model.forward([187, 510], None)
emb, state = model.forward([1563], state)           # RNN has state (use deepcopy to clone states)
emb, state = model.forward([310, 247, EOS_INDEX], state)
print(emb.detach().cpu().numpy())                   # same result as above

# batch inference (all sequences must share the same length)
batch_tokens = [
    [187, 510, 1563, 310],
    [247, EOS_INDEX, 187, 310],
]
emb_batch, batch_state = model.forward(batch_tokens, None, full_output=False)
print(emb_batch.detach().cpu().numpy())             # shape: [batch, n_vocab]
print(len(batch_state), batch_state[-2].shape)      # batched state shapes
print('\n')