AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-12 15:19:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix/jina tokenizer cache (#2735)

Browse Source
This commit is contained in:
Yeuoly 2024-03-07 16:32:37 +08:00 committed by GitHub
parent 1809f05904
commit 8fe83750b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 20 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
api/core/model_runtime/model_providers/jina/text_embedding/jina_tokenizer.py
View File

@ -1,20 +1,32 @@
from os.path import abspath, dirname, join from os.path import abspath, dirname, join
from threading import Lock
from transformers import AutoTokenizer from transformers import AutoTokenizer
class JinaTokenizer: class JinaTokenizer:
@staticmethod _tokenizer = None
def _get_num_tokens_by_jina_base(text: str) -> int: _lock = Lock()
@classmethod
def _get_tokenizer(cls):
if cls._tokenizer is None:
with cls._lock:
if cls._tokenizer is None:
base_path = abspath(__file__)
gpt2_tokenizer_path = join(dirname(base_path), 'tokenizer')
cls._tokenizer = AutoTokenizer.from_pretrained(gpt2_tokenizer_path)
return cls._tokenizer
@classmethod
def _get_num_tokens_by_jina_base(cls, text: str) -> int:
""" """
use jina tokenizer to get num tokens use jina tokenizer to get num tokens
""" """
base_path = abspath(__file__) tokenizer = cls._get_tokenizer()
gpt2_tokenizer_path = join(dirname(base_path), 'tokenizer')
tokenizer = AutoTokenizer.from_pretrained(gpt2_tokenizer_path)
tokens = tokenizer.encode(text) tokens = tokenizer.encode(text)
return len(tokens) return len(tokens)
@staticmethod @classmethod
def get_num_tokens(text: str) -> int: def get_num_tokens(cls, text: str) -> int:
return JinaTokenizer._get_num_tokens_by_jina_base(text) return cls._get_num_tokens_by_jina_base(text)