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