mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-07-09 21:19:20 +08:00
37 lines
1.0 KiB
Python
37 lines
1.0 KiB
Python
"""
|
|
Module for counting tokens and splitting text into chunks
|
|
"""
|
|
|
|
from langchain_core.language_models.chat_models import BaseChatModel
|
|
from langchain_mistralai import ChatMistralAI
|
|
from langchain_ollama import ChatOllama
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
|
|
def num_tokens_calculus(string: str, llm_model: BaseChatModel) -> int:
|
|
"""
|
|
Returns the number of tokens in a text string.
|
|
"""
|
|
if isinstance(llm_model, ChatOpenAI):
|
|
from .tokenizers.tokenizer_openai import num_tokens_openai
|
|
|
|
num_tokens_fn = num_tokens_openai
|
|
|
|
elif isinstance(llm_model, ChatMistralAI):
|
|
from .tokenizers.tokenizer_mistral import num_tokens_mistral
|
|
|
|
num_tokens_fn = num_tokens_mistral
|
|
|
|
elif isinstance(llm_model, ChatOllama):
|
|
from .tokenizers.tokenizer_ollama import num_tokens_ollama
|
|
|
|
num_tokens_fn = num_tokens_ollama
|
|
|
|
else:
|
|
from .tokenizers.tokenizer_openai import num_tokens_openai
|
|
|
|
num_tokens_fn = num_tokens_openai
|
|
|
|
num_tokens = num_tokens_fn(string, llm_model)
|
|
return num_tokens
|