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

Merge pull request #877 from ScrapeGraphAI/856-token-indices-sequence-length

Browse Source 
856 token indices sequence length
This commit is contained in:
Marco Perini 2025-01-12 12:48:36 +01:00 committed by GitHub
commit a523df0c80
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 55 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
examples/local_models/smart_scraper_ollama.py
View File

@ -15,7 +15,7 @@ graph_config = {
"temperature": 0,
"format": "json", # Ollama needs the format to be specified explicitly
# "base_url": "http://localhost:11434", # set ollama URL arbitrarily
"model_tokens": 1024,
"model_tokens": 4096,
},
"verbose": True,
"headless": False,
@ -25,7 +25,7 @@ graph_config = {
# Create the SmartScraperGraph instance and run it
# ************************************************
smart_scraper_graph = SmartScraperGraph(
prompt="Find some information about what does the company do, the name and a contact email.",
prompt="Find some information about what does the company do and the list of founders.",
source="https://scrapegraphai.com/",
config=graph_config,
)

29
examples/local_models/smart_scraper_schema_ollama.py
View File

@ -1,12 +1,15 @@
"""
"""
Basic example of scraping pipeline using SmartScraper with schema
"""
import json
from typing import List
from pydantic import BaseModel, Field
from scrapegraphai.graphs import SmartScraperGraph
from scrapegraphai.utils import prettify_exec_info
# ************************************************
# Define the configuration for the graph
# ************************************************
@ -14,18 +17,15 @@ class Project(BaseModel):
title: str = Field(description="The title of the project")
description: str = Field(description="The description of the project")
class Projects(BaseModel):
projects: List[Project]
projects: list[Project]
graph_config = {
"llm": {
"model": "ollama/llama3.1",
"temperature": 0,
"format": "json", # Ollama needs the format to be specified explicitly
# "base_url": "http://localhost:11434", # set ollama URL arbitrarily
},
"llm": {"model": "ollama/llama3.2", "temperature": 0, "model_tokens": 4096},
"verbose": True,
"headless": False
"headless": False,
}
# ************************************************
@ -36,8 +36,15 @@ smart_scraper_graph = SmartScraperGraph(
prompt="List me all the projects with their description",
source="https://perinim.github.io/projects/",
schema=Projects,
config=graph_config
config=graph_config,
)
result = smart_scraper_graph.run()
print(json.dumps(result, indent=4))
# ************************************************
# Get graph execution info
# ************************************************
graph_exec_info = smart_scraper_graph.get_execution_info()
print(prettify_exec_info(graph_exec_info))

3
pyproject.toml
View File

@ -30,8 +30,7 @@ dependencies = [
"googlesearch-python>=1.2.5",
"async-timeout>=4.0.3",
"simpleeval>=1.0.0",
"jsonschema>=4.23.0",
"transformers>=4.46.3",
"jsonschema>=4.23.0"
]
readme = "README.md"

4
scrapegraphai/docloaders/chromium.py
View File

@ -61,7 +61,6 @@ class ChromiumLoader(BaseLoader):
dynamic_import(backend, message)
self.backend = backend
self.browser_config = kwargs
self.headless = headless
self.proxy = parse_or_search_proxy(proxy) if proxy else None
@ -69,7 +68,8 @@ class ChromiumLoader(BaseLoader):
self.load_state = load_state
self.requires_js_support = requires_js_support
self.storage_state = storage_state
self.browser_name = browser_name
self.backend = kwargs.get("backend", backend)
self.browser_name = kwargs.get("browser_name", browser_name)
self.retry_limit = kwargs.get("retry_limit", retry_limit)
self.timeout = kwargs.get("timeout", timeout)

5
scrapegraphai/graphs/abstract_graph.py
View File

@ -203,8 +203,9 @@ class AbstractGraph(ABC):
]
except KeyError:
print(
f"""Model {llm_params['model_provider']}/{llm_params['model']} not found,
using default token size (8192)"""
f"""Max input tokens for model {llm_params['model_provider']}/{llm_params['model']} not found,
please specify the model_tokens parameter in the llm section of the graph configuration.
Using default token size: 8192"""
)
self.model_token = 8192
else:

10
scrapegraphai/nodes/generate_answer_node.py
View File

@ -10,7 +10,7 @@ from langchain_aws import ChatBedrock
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from langchain_openai import AzureChatOpenAI, ChatOpenAI
from langchain_openai import ChatOpenAI
from requests.exceptions import Timeout
from tqdm import tqdm
@ -59,7 +59,10 @@ class GenerateAnswerNode(BaseNode):
self.llm_model = node_config["llm_model"]
if isinstance(node_config["llm_model"], ChatOllama):
self.llm_model.format = "json"
if node_config.get("schema", None) is None:
self.llm_model.format = "json"
else:
self.llm_model.format = self.node_config["schema"].model_json_schema()
self.verbose = node_config.get("verbose", False)
self.force = node_config.get("force", False)
@ -123,8 +126,7 @@ class GenerateAnswerNode(BaseNode):
format_instructions = ""
if (
isinstance(self.llm_model, (ChatOpenAI, AzureChatOpenAI))
and not self.script_creator
not self.script_creator
or self.force
and not self.script_creator
or self.is_md_scraper

13
scrapegraphai/nodes/generate_answer_node_k_level.py
View File

@ -6,10 +6,11 @@ from typing import List, Optional
from langchain.prompts import PromptTemplate
from langchain_aws import ChatBedrock
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from langchain_mistralai import ChatMistralAI
from langchain_openai import AzureChatOpenAI, ChatOpenAI
from langchain_openai import ChatOpenAI
from tqdm import tqdm
from ..prompts import (
@ -55,6 +56,13 @@ class GenerateAnswerNodeKLevel(BaseNode):
super().__init__(node_name, "node", input, output, 2, node_config)
self.llm_model = node_config["llm_model"]
if isinstance(node_config["llm_model"], ChatOllama):
if node_config.get("schema", None) is None:
self.llm_model.format = "json"
else:
self.llm_model.format = self.node_config["schema"].model_json_schema()
self.embedder_model = node_config.get("embedder_model", None)
self.verbose = node_config.get("verbose", False)
self.force = node_config.get("force", False)
@ -92,8 +100,7 @@ class GenerateAnswerNodeKLevel(BaseNode):
format_instructions = ""
if (
isinstance(self.llm_model, (ChatOpenAI, AzureChatOpenAI))
and not self.script_creator
not self.script_creator
or self.force
and not self.script_creator
or self.is_md_scraper

4
scrapegraphai/nodes/parse_node.py
View File

@ -96,7 +96,6 @@ class ParseNode(BaseNode):
chunks = split_text_into_chunks(
text=docs_transformed.page_content,
chunk_size=self.chunk_size - 250,
model=self.llm_model,
)
else:
docs_transformed = docs_transformed[0]
@ -115,11 +114,10 @@ class ParseNode(BaseNode):
chunks = split_text_into_chunks(
text=docs_transformed.page_content,
chunk_size=chunk_size,
model=self.llm_model,
)
else:
chunks = split_text_into_chunks(
text=docs_transformed, chunk_size=chunk_size, model=self.llm_model
text=docs_transformed, chunk_size=chunk_size
)
state.update({self.output[0]: chunks})

14
scrapegraphai/utils/split_text_into_chunks.py
View File

@ -4,14 +4,10 @@ split_text_into_chunks module
from typing import List
from langchain_core.language_models.chat_models import BaseChatModel
from .tokenizer import num_tokens_calculus
def split_text_into_chunks(
text: str, chunk_size: int, model: BaseChatModel, use_semchunk=True
) -> List[str]:
def split_text_into_chunks(text: str, chunk_size: int, use_semchunk=True) -> List[str]:
"""
Splits the text into chunks based on the number of tokens.
@ -27,9 +23,9 @@ def split_text_into_chunks(
from semchunk import chunk
def count_tokens(text):
return num_tokens_calculus(text, model)
return num_tokens_calculus(text)
chunk_size = min(chunk_size - 500, int(chunk_size * 0.9))
chunk_size = min(chunk_size, int(chunk_size * 0.9))
chunks = chunk(
text=text, chunk_size=chunk_size, token_counter=count_tokens, memoize=False
@ -37,7 +33,7 @@ def split_text_into_chunks(
return chunks
else:
tokens = num_tokens_calculus(text, model)
tokens = num_tokens_calculus(text)
if tokens <= chunk_size:
return [text]
@ -48,7 +44,7 @@ def split_text_into_chunks(
words = text.split()
for word in words:
word_tokens = num_tokens_calculus(word, model)
word_tokens = num_tokens_calculus(word)
if current_length + word_tokens > chunk_size:
chunks.append(" ".join(current_chunk))
current_chunk = [word]

28
scrapegraphai/utils/tokenizer.py
View File

@ -2,35 +2,15 @@
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
from .tokenizers.tokenizer_openai import num_tokens_openai
def num_tokens_calculus(string: str, llm_model: BaseChatModel) -> int:
def num_tokens_calculus(string: str) -> 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
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)
num_tokens = num_tokens_fn(string)
return num_tokens

6
scrapegraphai/utils/tokenizers/tokenizer_openai.py
View File

@ -3,19 +3,17 @@ Tokenization utilities for OpenAI models
"""
import tiktoken
from langchain_core.language_models.chat_models import BaseChatModel
from ..logging import get_logger
def num_tokens_openai(text: str, llm_model: BaseChatModel) -> int:
def num_tokens_openai(text: str) -> int:
"""
Estimate the number of tokens in a given text using OpenAI's tokenization method,
adjusted for different OpenAI models.
Args:
text (str): The text to be tokenized and counted.
llm_model (BaseChatModel): The specific OpenAI model to adjust tokenization.
Returns:
int: The number of tokens in the text.
@ -25,7 +23,7 @@ def num_tokens_openai(text: str, llm_model: BaseChatModel) -> int:
logger.debug(f"Counting tokens for text of {len(text)} characters")
encoding = tiktoken.encoding_for_model("gpt-4")
encoding = tiktoken.encoding_for_model("gpt-4o")
num_tokens = len(encoding.encode(text))
return num_tokens

4
uv.lock
View File

@ -3429,7 +3429,7 @@ wheels = [
[[package]]
name = "scrapegraphai"
version = "1.35.0b2"
version = "1.35.0"
source = { editable = "." }
dependencies = [
{ name = "async-timeout", version = "4.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
@ -3452,7 +3452,6 @@ dependencies = [
{ name = "simpleeval" },
{ name = "tiktoken" },
{ name = "tqdm" },
{ name = "transformers" },
{ name = "undetected-playwright" },
]
@ -3516,7 +3515,6 @@ requires-dist = [
{ name = "surya-ocr", marker = "extra == 'ocr'", specifier = ">=0.5.0" },
{ name = "tiktoken", specifier = ">=0.7" },
{ name = "tqdm", specifier = ">=4.66.4" },
{ name = "transformers", specifier = ">=4.46.3" },
{ name = "undetected-playwright", specifier = ">=0.3.0" },
]