harrisonvanderbyl / rwkvstic

Framework agnostic python runtime for RWKV models
https://hazzzardous-rwkv-instruct.hf.space
MIT License
144 stars 18 forks source link

Misguidance of onnx export in the README.md #21

Closed L-M-Sherlock closed 1 year ago

L-M-Sherlock commented 1 year ago

I find that there is not kwarg named backend in the RWKV(), so I replace backend with mode.

Before:

from rwkvstic.load import RWKV
from rwkvstic.agnostic.backends import ONNX_EXPORT
import torch

model = RWKV("model.pth", backend=ONNX_EXPORT, dtype=torch.float16) # or torch.float32
# the model is exported to model_{layers}_{embed}.onnx
# the external data is stored in model_{layers}_{embed}.bin

After

from rwkvstic.load import RWKV
from rwkvstic.agnostic.backends import ONNX_EXPORT
import torch

model = RWKV("model.pth", mode=ONNX_EXPORT, dtype=torch.float16) # or torch.float32
# the model is exported to model_{layers}_{embed}.onnx
# the external data is stored in model_{layers}_{embed}.bin