fix: model count

This commit is contained in:
Marco Vinciguerra 2024-08-16 17:38:55 +02:00
parent 6730797008
commit faef3186f7
4 changed files with 4 additions and 2 deletions

View File

@ -9,7 +9,7 @@ from scrapegraphai.utils import prettify_exec_info
graph_config = {
"llm": {
"model": "ollama/llama3.1",
"model": "ollama/mistral",
"temperature": 0,
"format": "json", # Ollama needs the format to be specified explicitly
# "base_url": "http://localhost:11434", # set ollama URL arbitrarily

View File

@ -149,7 +149,7 @@ class AbstractGraph(ABC):
known_models = ["openai", "azure_openai", "google_genai", "ollama", "oneapi", "nvidia", "groq", "google_vertexai", "bedrock", "mistralai", "hugging_face", "deepseek", "ernie", "fireworks"]
if llm_params["model"] not in known_models:
if llm_params["model"].split("/")[0] not in known_models:
raise ValueError(f"Model '{llm_params['model']}' is not supported")
try:

View File

@ -62,6 +62,7 @@ models_tokens = {
"scrapegraph": 8192,
"llava": 4096,
"mixtral:8x22b-instruct": 65536,
"mistral":8192,
"mistral-openorca": 32000,
"nomic-embed-text": 8192,
"nous-hermes2:34b": 4096,

View File

@ -80,6 +80,7 @@ class ParseNode(BaseNode):
docs_transformed = docs_transformed[0]
if isinstance(docs_transformed, Document):
chunks = chunk(text=docs_transformed.page_content,
chunk_size=self.node_config.get("chunk_size", 4096)-250,
token_counter=lambda text: len(text.split()),