From d32351c1302e9f91d9345addfcc1bcb8830516c4 Mon Sep 17 00:00:00 2001 From: josc146 Date: Mon, 19 Jun 2023 22:30:49 +0800 Subject: [PATCH] exact model name --- backend-python/routes/completion.py | 6 +++--- backend-python/utils/rwkv.py | 2 ++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/backend-python/routes/completion.py b/backend-python/routes/completion.py index e572e65..841028d 100644 --- a/backend-python/routes/completion.py +++ b/backend-python/routes/completion.py @@ -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": { diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py index 0453adc..208a961 100644 --- a/backend-python/utils/rwkv.py +++ b/backend-python/utils/rwkv.py @@ -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