zhuchaowe/maxkb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: VLLM supplier recalculates token function (#2375)

Browse Source
This commit is contained in:
shaohuzhang1 2025-02-24 17:15:46 +08:00 committed by GitHub
parent 2d4deda6b4
commit 6b72611b72
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 32 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
apps/setting/models_provider/impl/vllm_model_provider/model/image.py
View File

@ -1,5 +1,8 @@
from typing import Dict
from typing import Dict, List
from langchain_core.messages import get_buffer_string, BaseMessage
from common.config.tokenizer_manage_config import TokenizerManage
from setting.models_provider.base_model_provider import MaxKBBaseModel
from setting.models_provider.impl.base_chat_open_ai import BaseChatOpenAI
@ -21,3 +24,15 @@ class VllmImage(MaxKBBaseModel, BaseChatOpenAI):
def is_cache_model(self):
return False
def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
if self.usage_metadata is None or self.usage_metadata == {}:
tokenizer = TokenizerManage.get_tokenizer()
return sum([len(tokenizer.encode(get_buffer_string([m]))) for m in messages])
return self.usage_metadata.get('input_tokens', 0)
def get_num_tokens(self, text: str) -> int:
if self.usage_metadata is None or self.usage_metadata == {}:
tokenizer = TokenizerManage.get_tokenizer()
return len(tokenizer.encode(text))
return self.get_last_generation_info().get('output_tokens', 0)

17
apps/setting/models_provider/impl/vllm_model_provider/model/llm.py
View File

@ -1,8 +1,11 @@
# coding=utf-8
from typing import Dict
from typing import Dict, List
from urllib.parse import urlparse, ParseResult
from langchain_core.messages import BaseMessage, get_buffer_string
from common.config.tokenizer_manage_config import TokenizerManage
from setting.models_provider.base_model_provider import MaxKBBaseModel
from setting.models_provider.impl.base_chat_open_ai import BaseChatOpenAI
@ -33,3 +36,15 @@ class VllmChatModel(MaxKBBaseModel, BaseChatOpenAI):
stream_usage=True,
)
return vllm_chat_open_ai
def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
if self.usage_metadata is None or self.usage_metadata == {}:
tokenizer = TokenizerManage.get_tokenizer()
return sum([len(tokenizer.encode(get_buffer_string([m]))) for m in messages])
return self.usage_metadata.get('input_tokens', 0)
def get_num_tokens(self, text: str) -> int:
if self.usage_metadata is None or self.usage_metadata == {}:
tokenizer = TokenizerManage.get_tokenizer()
return len(tokenizer.encode(text))
return self.get_last_generation_info().get('output_tokens', 0)