Skip to content

Commit b6f1766

Browse files
VinciGit00f-aguzziwangdongpeng1
committed
add OneAPI integration
Co-Authored-By: Federico Aguzzi <62149513+f-aguzzi@users.noreply.github.com> Co-Authored-By: wangdongpeng1 <74647183+wangdongpeng1@users.noreply.github.com>
1 parent e1006f3 commit b6f1766

File tree

4 files changed

+75
-6
lines changed

4 files changed

+75
-6
lines changed
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
"""
2+
Basic example of scraping pipeline using SmartScraper
3+
"""
4+
5+
from scrapegraphai.graphs import SmartScraperGraph
6+
from scrapegraphai.utils import prettify_exec_info
7+
8+
# ************************************************
9+
# Define the configuration for the graph
10+
# *********************************************
11+
12+
graph_config = {
13+
"llm": {
14+
"api_key": "***************************",
15+
"model": "oneapi/qwen-turbo",
16+
"base_url": "http://127.0.0.1:3000/v1", # 设置 OneAPI URL
17+
},
18+
"embeddings": {
19+
"model": "ollama/nomic-embed-text",
20+
"base_url": "http://127.0.0.1:11434", # 设置 Ollama URL
21+
}
22+
}
23+
24+
# ************************************************
25+
# Create the SmartScraperGraph instance and run it
26+
# ************************************************
27+
28+
smart_scraper_graph = SmartScraperGraph(
29+
prompt="该网站为XXXXX,请提取出标题、发布时间、发布来源以及内容摘要,并以中文回答。",
30+
# 也可以使用已下载的 HTML 代码的字符串
31+
source="http://XXXX",
32+
config=graph_config
33+
)
34+
35+
# ************************************************
36+
# Get graph execution info
37+
# ************************************************
38+
result = smart_scraper_graph.run()
39+
print(result)
40+
print(prettify_exec_info(result))

scrapegraphai/graphs/abstract_graph.py

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
HuggingFace,
2222
Ollama,
2323
OpenAI,
24+
OneApi
2425
)
2526
from ..utils.logging import set_verbosity_debug, set_verbosity_warning
2627

@@ -54,19 +55,20 @@ class AbstractGraph(ABC):
5455
... # Implementation of graph creation here
5556
... return graph
5657
...
57-
>>> my_graph = MyGraph("Example Graph", {"llm": {"model": "gpt-3.5-turbo"}}, "example_source")
58+
>>> my_graph = MyGraph("Example Graph",
59+
{"llm": {"model": "gpt-3.5-turbo"}}, "example_source")
5860
>>> result = my_graph.run()
5961
"""
6062

61-
def __init__(self, prompt: str, config: dict, source: Optional[str] = None, schema: Optional[str] = None):
63+
def __init__(self, prompt: str, config: dict,
64+
source: Optional[str] = None, schema: Optional[str] = None):
6265

6366
self.prompt = prompt
6467
self.source = source
6568
self.config = config
6669
self.schema = schema
6770
self.llm_model = self._create_llm(config["llm"], chat=True)
68-
self.embedder_model = self._create_default_embedder(llm_config=config["llm"]
69-
) if "embeddings" not in config else self._create_embedder(
71+
self.embedder_model = self._create_default_embedder(llm_config=config["llm"] ) if "embeddings" not in config else self._create_embedder(
7072
config["embeddings"])
7173
self.verbose = False if config is None else config.get(
7274
"verbose", False)
@@ -98,7 +100,7 @@ def __init__(self, prompt: str, config: dict, source: Optional[str] = None, sche
98100
"llm_model": self.llm_model,
99101
"embedder_model": self.embedder_model
100102
}
101-
103+
102104
self.set_common_params(common_params, overwrite=False)
103105

104106
def set_common_params(self, params: dict, overwrite=False):
@@ -163,7 +165,14 @@ def _create_llm(self, llm_config: dict, chat=False) -> object:
163165
except KeyError as exc:
164166
raise KeyError("Model not supported") from exc
165167
return OpenAI(llm_params)
166-
168+
elif "oneapi" in llm_params["model"]:
169+
# take the model after the last dash
170+
llm_params["model"] = llm_params["model"].split("/")[-1]
171+
try:
172+
self.model_token = models_tokens["oneapi"][llm_params["model"]]
173+
except KeyError as exc:
174+
raise KeyError("Model Model not supported") from exc
175+
return OneApi(llm_params)
167176
elif "azure" in llm_params["model"]:
168177
# take the model after the last dash
169178
llm_params["model"] = llm_params["model"].split("/")[-1]

scrapegraphai/helpers/models_tokens.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,9 @@
8080
"snowflake-arctic-embed:l": 8192,
8181
"mxbai-embed-large": 512,
8282
},
83+
"oneapi": {
84+
"qwen-turbo": 16380
85+
},
8386
"groq": {
8487
"llama3-8b-8192": 8192,
8588
"llama3-70b-8192": 8192,

scrapegraphai/models/oneapi.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
"""
2+
OpenAI Module
3+
"""
4+
from langchain_openai import ChatOpenAI
5+
6+
7+
class OneApi(ChatOpenAI):
8+
"""
9+
A wrapper for the OneApi class that provides default configuration
10+
and could be extended with additional methods if needed.
11+
12+
Args:
13+
llm_config (dict): Configuration parameters for the language model.
14+
"""
15+
16+
def __init__(self, llm_config: dict):
17+
super().__init__(**llm_config)

0 commit comments

Comments
 (0)