https://github.com/cryscan/web-rwkv-py

https://github.com/cryscan/web-rwkv-py

Last synced: 2 months ago
JSON representation

• Host: GitHub
• URL: https://github.com/cryscan/web-rwkv-py
• Owner: cryscan
• Created: 2023-12-13T18:00:00.000Z (about 1 year ago)
• Default Branch: main
• Last Pushed: 2024-06-15T13:05:48.000Z (7 months ago)
• Last Synced: 2024-10-20T11:43:16.863Z (2 months ago)
• Language: Rust
• Size: 494 KB
• Stars: 9
• Watchers: 2
• Forks: 5
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

README

        
# Web-RWKV-Py

Python binding for [`web-rwkv`](https://github.com/cryscan/web-rwkv).

# Todos

- [x] Basic V5 inference support

- [x] Support V4, V5 and V6

- [ ] Batched inference

# Usage

1. Install python and rust.

2. Install maturin by

   

   ```bash

   $ pip install maturin

   ```

4. Build and install:

   ```bash

   $ maturin develop --release

   ```

5. Try using `web-rwkv` in python:

   ```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

   )

   model.clear_state()

   logits = model.run([114, 514])

   ```

   

# Advanced Usage

1. Get, clone and load current state:

   ```python

   logits = model.run([114, 514])

   state = model.back_state(wrp.StateDevice.Gpu)

   # state = model.back_state(wrp.StateDevice.Cpu)

   state_cloned = state.deep_clone()

   model.load_state(state_cloned)

   logits = model.run([1919, 810])

   ```

   

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

   ```python

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

   assert(len(logits) == 4)

   ```