CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-07-01 21:00:48 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix(AbstractGraph): LangChain warnings handling, Mistral tokens

Browse Source
This commit is contained in:
Federico Aguzzi 2024-08-07 15:15:35 +02:00
parent 07720b6e0c
commit 786af992f8
2 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
scrapegraphai/graphs/abstract_graph.py
View File

@ -5,6 +5,7 @@ AbstractGraph Module
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Optional from typing import Optional
import uuid import uuid
import warnings
from pydantic import BaseModel from pydantic import BaseModel
from langchain_community.chat_models import ErnieBotChat from langchain_community.chat_models import ErnieBotChat
@ -144,7 +145,9 @@ class AbstractGraph(ABC):
self.model_token = default_token self.model_token = default_token
llm_params["model_provider"] = provider llm_params["model_provider"] = provider
llm_params["model"] = model_name llm_params["model"] = model_name
return init_chat_model(**llm_params) with warnings.catch_warnings():
warnings.simplefilter("ignore")
return init_chat_model(**llm_params)
if "azure" in llm_params["model"]: if "azure" in llm_params["model"]:
model_name = llm_params["model"].split("/")[-1] model_name = llm_params["model"].split("/")[-1]

2
scrapegraphai/helpers/models_tokens.py
View File

@ -145,7 +145,7 @@ models_tokens = {
"cohere.embed-english-v3": 512, "cohere.embed-english-v3": 512,
"cohere.embed-multilingual-v3": 512 "cohere.embed-multilingual-v3": 512
}, },
"mistral": { "mistralai": {
"mistral-large-latest": 128000, "mistral-large-latest": 128000,
"open-mistral-nemo": 128000, "open-mistral-nemo": 128000,
"codestral-latest": 32000, "codestral-latest": 32000,