perf: Add custom token_ids method (#2264)

This commit is contained in:
shaohuzhang1 2025-02-13 15:40:09 +08:00 committed by GitHub
parent dac0a9c3f3
commit 051dc13322
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -15,8 +15,14 @@ from langchain_openai.chat_models.base import _convert_chunk_to_generation_chunk
from common.config.tokenizer_manage_config import TokenizerManage
def custom_get_token_ids(text: str):
tokenizer = TokenizerManage.get_tokenizer()
return tokenizer.encode(text)
class BaseChatOpenAI(ChatOpenAI):
usage_metadata: dict = {}
custom_get_token_ids = custom_get_token_ids
def get_last_generation_info(self) -> Optional[Dict[str, Any]]:
return self.usage_metadata