theluyuan/RWKV-Runner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

exact model name

Browse Source
This commit is contained in:
josc146 2023-06-19 22:30:49 +08:00
parent 967be6f88f
commit d32351c130
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
backend-python/routes/completion.py
View File

@ -117,7 +117,7 @@ async def eval_rwkv(
yield json.dumps(
{
"response": response,
"model": "rwkv",
"model": model.name,
"choices": [
{
"delta": {"content": delta},
@ -153,7 +153,7 @@ async def eval_rwkv(
yield json.dumps(
{
"response": response,
"model": "rwkv",
"model": model.name,
"choices": [
{
"delta": {},
@ -173,7 +173,7 @@ async def eval_rwkv(
else:
yield {
"response": response,
"model": "rwkv",
"model": model.name,
"choices": [
{
"message": {

2
backend-python/utils/rwkv.py
View File

@ -21,6 +21,8 @@ class RWKV:
def __init__(self, model: str, strategy: str, tokens_path: str) -> None:
from rwkv.model import RWKV as Model # dynamic import to make RWKV_CUDA_ON work
filename, _ = os.path.splitext(os.path.basename(model))
self.name = filename
self.model = Model(model, strategy)
self.pipeline = PIPELINE(self.model, tokens_path)
self.model_state = None