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( yield json.dumps(
{ {
"response": response, "response": response,
"model": "rwkv", "model": model.name,
"choices": [ "choices": [
{ {
"delta": {"content": delta}, "delta": {"content": delta},
@ -153,7 +153,7 @@ async def eval_rwkv(
yield json.dumps( yield json.dumps(
{ {
"response": response, "response": response,
"model": "rwkv", "model": model.name,
"choices": [ "choices": [
{ {
"delta": {}, "delta": {},
@ -173,7 +173,7 @@ async def eval_rwkv(
else: else:
yield { yield {
"response": response, "response": response,
"model": "rwkv", "model": model.name,
"choices": [ "choices": [
{ {
"message": { "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: 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 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.model = Model(model, strategy)
self.pipeline = PIPELINE(self.model, tokens_path) self.pipeline = PIPELINE(self.model, tokens_path)
self.model_state = None self.model_state = None