web-rwkv-py 0.2.0

No project description provided

Web-RWKV-Py

Python binding for web-rwkv.

Todos

• Basic V5 inference support
• Support V4, V5 and V6
• Batched inference

Usage

1. Install python and rust.

2. Install maturin by

   $ pip install maturin
   

3. Build and install:

   $ maturin develop --release
   

4. Try using web-rwkv in python:

   import web_rwkv_py as wrp
   
   model = wrp.Model(
      "/path/to/model.st", # model path
      quant=0,             # int8 quantization layers
      quant_nf4=0,         # nf4 quantization layers
   )
   logits, state = model.run([114, 514], state=None)
   

Advanced Usage

1. Clone states:

   logits, state = model.run([114, 514], state=None)
   state_cloned = model.clone_state(state)
   

2. Return predictions of all tokens (not only the last's):

   logits, state = model.run_full([114, 514, 1919, 810], state=None)
   assert(len(logits) == 4)
   

3. Specify token chunk size (default is 128, the larger the faster prefilling is, but you will need a beefy GPU):

   tokens = [114, 514, 1919, 810]
   logits, state = model.run_full(tokens, state=None, token_chunk_size=256)