style: fix ruff check
This commit is contained in:
Wenjie Zhang
parent
911f88d7eb
commit
a6857c3e21
@ -41,10 +41,10 @@ dependencies = [
|
||||
]
|
||||
[tool.ruff]
|
||||
line-length = 140 # 代码最大行宽
|
||||
select = [ # 选择的规则
|
||||
lint.select = [ # 选择的规则
|
||||
"F",
|
||||
"E",
|
||||
"W",
|
||||
"UP",
|
||||
]
|
||||
ignore = ["F401"] # 忽略的规则
|
||||
lint.ignore = ["F401"] # 忽略的规则
|
||||
@ -1,35 +0,0 @@
|
||||
dashscope>=1.20.3
|
||||
PyMuPDF==1.23.26
|
||||
peft>=0.11.1
|
||||
FlagEmbedding>=1.3.2
|
||||
Flask>=3.0.3
|
||||
Flask_Cors>=4.0.1
|
||||
llama_index>=0.11.8
|
||||
openai>=1.44.1
|
||||
paddleocr>=2.8.1
|
||||
paddlepaddle>=2.4.0
|
||||
pandas>=2.2.2
|
||||
Pillow>=10.4.0
|
||||
pymilvus>=2.4.4
|
||||
python-dotenv>=1.0.1
|
||||
PyYAML>=6.0.2
|
||||
qianfan>=0.4.7
|
||||
torch>=2.4.0
|
||||
tqdm>=4.66.1
|
||||
zhipuai>=2.1.2
|
||||
neo4j>=5.22.0
|
||||
sentencepiece>=0.1.99
|
||||
llama-index-readers-file
|
||||
opencv-python-headless
|
||||
docx2txt
|
||||
uvicorn[standard]
|
||||
fastapi
|
||||
python-multipart
|
||||
tavily-python
|
||||
rapidocr_onnxruntime
|
||||
langchain
|
||||
langsmith
|
||||
langgraph
|
||||
langchain-openai
|
||||
langchain-community
|
||||
PyJWT>=2.10.1
|
||||
@ -17,4 +17,4 @@ outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
@ -15,4 +15,4 @@ chat_response = client.chat.completions.create(
|
||||
{"role": "user", "content": "Tell me a joke."},
|
||||
]
|
||||
)
|
||||
print("Chat response:", chat_response)
|
||||
print("Chat response:", chat_response)
|
||||
|
||||
@ -65,4 +65,4 @@ class DBManager:
|
||||
session.close()
|
||||
|
||||
# 创建全局数据库管理器实例
|
||||
db_manager = DBManager()
|
||||
db_manager = DBManager()
|
||||
|
||||
@ -1,3 +1,3 @@
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
|
||||
Base = declarative_base()
|
||||
Base = declarative_base()
|
||||
|
||||
@ -103,4 +103,4 @@ class KnowledgeNode(Base):
|
||||
"start_char_idx": self.start_char_idx,
|
||||
"end_char_idx": self.end_char_idx,
|
||||
"metadata": self.meta_info or {} # 确保映射正确
|
||||
}
|
||||
}
|
||||
|
||||
@ -18,4 +18,4 @@ class Thread(Base):
|
||||
update_at = Column(DateTime, default=func.now(), onupdate=func.now(), comment="更新时间")
|
||||
|
||||
description = Column(String(255), nullable=True, comment="描述")
|
||||
status = Column(Integer, default=1, comment="状态")
|
||||
status = Column(Integer, default=1, comment="状态")
|
||||
|
||||
@ -52,4 +52,4 @@ class OperationLog(Base):
|
||||
"details": self.details,
|
||||
"ip_address": self.ip_address,
|
||||
"timestamp": self.timestamp.isoformat() if self.timestamp else None
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,9 +1,8 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from fastapi.security import OAuth2PasswordRequestForm
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Optional
|
||||
from sqlalchemy.orm import Session
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import datetime
|
||||
|
||||
from server.db_manager import db_manager
|
||||
from server.models.user_model import User, OperationLog
|
||||
@ -27,16 +26,16 @@ class UserCreate(BaseModel):
|
||||
role: str = "user"
|
||||
|
||||
class UserUpdate(BaseModel):
|
||||
username: Optional[str] = None
|
||||
password: Optional[str] = None
|
||||
role: Optional[str] = None
|
||||
username: str | None = None
|
||||
password: str | None = None
|
||||
role: str | None = None
|
||||
|
||||
class UserResponse(BaseModel):
|
||||
id: int
|
||||
username: str
|
||||
role: str
|
||||
created_at: str
|
||||
last_login: Optional[str] = None
|
||||
last_login: str | None = None
|
||||
|
||||
class InitializeAdmin(BaseModel):
|
||||
username: str
|
||||
@ -202,7 +201,7 @@ async def create_user(
|
||||
return new_user.to_dict()
|
||||
|
||||
# 路由：获取所有用户（管理员权限）
|
||||
@auth.get("/users", response_model=List[UserResponse])
|
||||
@auth.get("/users", response_model=list[UserResponse])
|
||||
async def read_users(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
@ -344,4 +343,4 @@ async def delete_user(
|
||||
db.delete(user)
|
||||
db.commit()
|
||||
|
||||
return {"success": True, "message": "用户已删除"}
|
||||
return {"success": True, "message": "用户已删除"}
|
||||
|
||||
@ -1,15 +1,13 @@
|
||||
from fastapi import Request, Body, Depends
|
||||
from fastapi import APIRouter
|
||||
from fastapi import Request, Body
|
||||
|
||||
base = APIRouter()
|
||||
|
||||
from src import config, retriever, knowledge_base, graph_base
|
||||
from src.utils import logger
|
||||
from server.utils.auth_middleware import get_admin_user, get_superadmin_user
|
||||
from server.models.user_model import User
|
||||
|
||||
|
||||
base = APIRouter()
|
||||
|
||||
@base.get("/")
|
||||
async def route_index():
|
||||
return {"message": "You Got It!"}
|
||||
@ -49,7 +47,7 @@ def get_log(current_user: User = Depends(get_admin_user)):
|
||||
from src.utils.logging_config import LOG_FILE
|
||||
from collections import deque
|
||||
|
||||
with open(LOG_FILE, 'r') as f:
|
||||
with open(LOG_FILE) as f:
|
||||
last_lines = deque(f, maxlen=1000)
|
||||
|
||||
log = ''.join(last_lines)
|
||||
|
||||
@ -8,7 +8,6 @@ from fastapi import APIRouter, Body, Depends, HTTPException, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
from langchain_core.messages import AIMessageChunk, HumanMessage
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel
|
||||
|
||||
from src import executor, config, retriever
|
||||
@ -270,7 +269,7 @@ async def save_agent_config(
|
||||
if result:
|
||||
return {"success": True, "message": f"智能体 {agent_name} 配置已保存"}
|
||||
else:
|
||||
raise HTTPException(status_code=500, detail=f"保存智能体配置失败")
|
||||
raise HTTPException(status_code=500, detail="保存智能体配置失败")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"保存智能体配置出错: {e}, {traceback.format_exc()}")
|
||||
@ -318,18 +317,18 @@ async def get_agent_config(
|
||||
# ==================== 线程管理 API ====================
|
||||
|
||||
class ThreadCreate(BaseModel):
|
||||
title: Optional[str] = None
|
||||
title: str | None = None
|
||||
agent_id: str
|
||||
description: Optional[str] = None
|
||||
metadata: Optional[dict] = None
|
||||
description: str | None = None
|
||||
metadata: dict | None = None
|
||||
|
||||
|
||||
class ThreadResponse(BaseModel):
|
||||
id: str
|
||||
user_id: str
|
||||
agent_id: str
|
||||
title: Optional[str]
|
||||
description: Optional[str]
|
||||
title: str | None = None
|
||||
description: str | None = None
|
||||
create_at: str
|
||||
update_at: str
|
||||
|
||||
@ -366,9 +365,9 @@ async def create_thread(
|
||||
}
|
||||
|
||||
|
||||
@chat.get("/threads", response_model=List[ThreadResponse])
|
||||
@chat.get("/threads", response_model=list[ThreadResponse])
|
||||
async def list_threads(
|
||||
agent_id: Optional[str] = None,
|
||||
agent_id: str | None = None,
|
||||
db: Session = Depends(get_db),
|
||||
current_user: User = Depends(get_required_user)
|
||||
):
|
||||
@ -420,8 +419,8 @@ async def delete_thread(
|
||||
|
||||
|
||||
class ThreadUpdate(BaseModel):
|
||||
title: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
title: str | None = None
|
||||
description: str | None = None
|
||||
|
||||
|
||||
@chat.put("/thread/{thread_id}", response_model=ThreadResponse)
|
||||
|
||||
@ -1,7 +1,6 @@
|
||||
import os
|
||||
import asyncio
|
||||
import traceback
|
||||
from typing import List, Optional
|
||||
from fastapi import APIRouter, File, UploadFile, HTTPException, Depends, Body, Form, Query
|
||||
|
||||
from src.utils import logger, hashstr
|
||||
@ -25,7 +24,7 @@ async def get_databases(current_user: User = Depends(get_admin_user)):
|
||||
async def create_database(
|
||||
database_name: str = Body(...),
|
||||
description: str = Body(...),
|
||||
dimension: Optional[int] = Body(None),
|
||||
dimension: int | None = Body(None),
|
||||
current_user: User = Depends(get_admin_user)
|
||||
):
|
||||
logger.debug(f"Create database {database_name}")
|
||||
@ -53,19 +52,19 @@ async def query_test(query: str = Body(...), meta: dict = Body(...), current_use
|
||||
return result
|
||||
|
||||
@data.post("/file-to-chunk")
|
||||
async def file_to_chunk(files: List[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||||
async def file_to_chunk(files: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||||
logger.debug(f"File to chunk: {files} {params=}")
|
||||
result = await knowledge_base.file_to_chunk(files, params=params)
|
||||
return result
|
||||
|
||||
@data.post("/url-to-chunk")
|
||||
async def url_to_chunk(urls: List[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||||
async def url_to_chunk(urls: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||||
logger.debug(f"Url to chunk: {urls}")
|
||||
result = await knowledge_base.url_to_chunk(urls, params=params)
|
||||
return result
|
||||
|
||||
@data.post("/add-by-file")
|
||||
async def create_document_by_file(db_id: str = Body(...), files: List[str] = Body(...), current_user: User = Depends(get_admin_user)):
|
||||
async def create_document_by_file(db_id: str = Body(...), files: list[str] = Body(...), current_user: User = Depends(get_admin_user)):
|
||||
logger.debug(f"Add document in {db_id} by file: {files}")
|
||||
try:
|
||||
await knowledge_base.add_files(db_id, files)
|
||||
@ -113,7 +112,7 @@ async def get_document_info(db_id: str, file_id: str, current_user: User = Depen
|
||||
@data.post("/upload")
|
||||
async def upload_file(
|
||||
file: UploadFile = File(...),
|
||||
db_id: Optional[str] = Query(None),
|
||||
db_id: str | None = Query(None),
|
||||
current_user: User = Depends(get_admin_user)
|
||||
):
|
||||
if not file.filename:
|
||||
@ -170,7 +169,7 @@ async def get_graph_nodes(kgdb_name: str, num: int, current_user: User = Depends
|
||||
return {"result": graph_base.format_general_results(result), "message": "success"}
|
||||
|
||||
@data.post("/graph/add-by-jsonl")
|
||||
async def add_graph_entity(file_path: str = Body(...), kgdb_name: Optional[str] = Body(None), current_user: User = Depends(get_admin_user)):
|
||||
async def add_graph_entity(file_path: str = Body(...), kgdb_name: str | None = Body(None), current_user: User = Depends(get_admin_user)):
|
||||
if not config.enable_knowledge_graph:
|
||||
return {"message": "知识图谱未启用", "status": "failed"}
|
||||
|
||||
|
||||
@ -1 +1 @@
|
||||
# utils包初始化文件
|
||||
# utils包初始化文件
|
||||
|
||||
@ -1,5 +1,4 @@
|
||||
from typing import Optional, List, Callable
|
||||
from fastapi import Depends, HTTPException, status, APIRouter
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from fastapi.security import OAuth2PasswordBearer
|
||||
from sqlalchemy.orm import Session
|
||||
from jose import JWTError, jwt
|
||||
@ -30,7 +29,7 @@ def get_db():
|
||||
db.close()
|
||||
|
||||
# 获取当前用户
|
||||
async def get_current_user(token: Optional[str] = Depends(oauth2_scheme), db: Session = Depends(get_db)):
|
||||
async def get_current_user(token: str | None = Depends(oauth2_scheme), db: Session = Depends(get_db)):
|
||||
credentials_exception = HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="无效的凭证",
|
||||
@ -66,7 +65,7 @@ async def get_current_user(token: Optional[str] = Depends(oauth2_scheme), db: Se
|
||||
return user
|
||||
|
||||
# 获取已登录用户（抛出401如果未登录）
|
||||
async def get_required_user(user: Optional[User] = Depends(get_current_user)):
|
||||
async def get_required_user(user: User | None = Depends(get_current_user)):
|
||||
if user is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
|
||||
@ -2,17 +2,17 @@ from dotenv import load_dotenv
|
||||
|
||||
load_dotenv("src/.env")
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures import ThreadPoolExecutor # noqa: E402
|
||||
executor = ThreadPoolExecutor()
|
||||
|
||||
from src.config import Config
|
||||
from src.config import Config # noqa: E402
|
||||
config = Config()
|
||||
|
||||
from src.core import KnowledgeBase
|
||||
from src.core import KnowledgeBase # noqa: E402
|
||||
knowledge_base = KnowledgeBase()
|
||||
|
||||
from src.core import GraphDatabase
|
||||
from src.core import GraphDatabase # noqa: E402
|
||||
graph_base = GraphDatabase()
|
||||
|
||||
from src.core.retriever import Retriever
|
||||
retriever = Retriever()
|
||||
from src.core.retriever import Retriever # noqa: E402
|
||||
retriever = Retriever()
|
||||
|
||||
@ -39,4 +39,4 @@ __all__ = ["agent_manager"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
pass
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
from .graph import ChatbotAgent
|
||||
from .configuration import ChatbotConfiguration
|
||||
|
||||
__all__ = ["ChatbotAgent", "ChatbotConfiguration"]
|
||||
__all__ = ["ChatbotAgent", "ChatbotConfiguration"]
|
||||
|
||||
@ -63,7 +63,6 @@ class ChatbotAgent(BaseAgent):
|
||||
if self.graph:
|
||||
return self.graph
|
||||
|
||||
conf = self.config_schema.from_runnable_config(config_schema, agent_name=self.name)
|
||||
workflow = StateGraph(State, config_schema=self.config_schema)
|
||||
workflow.add_node("chatbot", self.llm_call)
|
||||
workflow.add_node("tools", ToolNode(tools=list(get_all_tools().values())))
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
from .graph import ReActAgent
|
||||
from .configuration import ReActConfiguration
|
||||
|
||||
__all__ = ["ReActAgent", "ReActConfiguration"]
|
||||
__all__ = ["ReActAgent", "ReActConfiguration"]
|
||||
|
||||
@ -2,24 +2,12 @@ import os
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
|
||||
model = ChatOpenAI(model="glm-4-plus",
|
||||
api_key=os.getenv("ZHIPUAI_API_KEY"),
|
||||
base_url="https://open.bigmodel.cn/api/paas/v4/",
|
||||
temperature=0)
|
||||
|
||||
|
||||
# For this tutorial we will use custom tool that returns pre-defined values for weather in two cities (NYC & SF)
|
||||
|
||||
from typing import Literal, Annotated
|
||||
|
||||
from langchain_core.tools import tool, StructuredTool
|
||||
|
||||
|
||||
|
||||
tools = []
|
||||
|
||||
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
|
||||
graph = create_react_agent(model, tools=tools, checkpointer=InMemorySaver())
|
||||
|
||||
@ -4,8 +4,7 @@ import os
|
||||
import yaml
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Type, Annotated, Optional, TypedDict
|
||||
from enum import Enum
|
||||
from typing import Annotated, TypedDict
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass, fields, field
|
||||
|
||||
@ -33,7 +32,7 @@ class Configuration(dict):
|
||||
|
||||
@classmethod
|
||||
def from_runnable_config(
|
||||
cls, config: Optional[RunnableConfig] = None, agent_name: str = None
|
||||
cls, config: RunnableConfig | None = None, agent_name: str | None = None
|
||||
) -> Configuration:
|
||||
"""Create a Configuration instance from a RunnableConfig object.
|
||||
|
||||
@ -58,18 +57,18 @@ class Configuration(dict):
|
||||
|
||||
# 合并三级配置，注意优先级
|
||||
merged_config = {}
|
||||
for field in _fields:
|
||||
for config_field in _fields:
|
||||
# 1. 默认使用类默认值
|
||||
if hasattr(instance, field):
|
||||
merged_config[field] = getattr(instance, field)
|
||||
if hasattr(instance, config_field):
|
||||
merged_config[config_field] = getattr(instance, config_field)
|
||||
|
||||
# 2. 如果文件配置中有此字段，则覆盖
|
||||
if field in file_config:
|
||||
merged_config[field] = file_config[field]
|
||||
if config_field in file_config:
|
||||
merged_config[config_field] = file_config[config_field]
|
||||
|
||||
# 3. 如果运行时配置中有此字段，则覆盖
|
||||
if field in configurable:
|
||||
merged_config[field] = configurable[field]
|
||||
if config_field in configurable:
|
||||
merged_config[config_field] = configurable[config_field]
|
||||
|
||||
# 创建并返回配置实例
|
||||
# logger.debug(f"合并配置: {merged_config}")
|
||||
@ -82,7 +81,7 @@ class Configuration(dict):
|
||||
file_config = {}
|
||||
if os.path.exists(config_file_path):
|
||||
try:
|
||||
with open(config_file_path, 'r', encoding='utf-8') as f:
|
||||
with open(config_file_path, encoding='utf-8') as f:
|
||||
file_config = yaml.safe_load(f) or {}
|
||||
# logger.info(f"从文件加载智能体 {agent_name} 配置: {file_config}")
|
||||
except Exception as e:
|
||||
@ -158,7 +157,7 @@ class Configuration(dict):
|
||||
},
|
||||
)
|
||||
|
||||
class BaseAgent():
|
||||
class BaseAgent:
|
||||
|
||||
"""
|
||||
定义一个基础 Agent 供 各类 graph 继承
|
||||
@ -245,4 +244,4 @@ class BaseAgent():
|
||||
必须确保在编译时设置 checkpointer，否则将无法获取历史记录。
|
||||
例如: graph = workflow.compile(checkpointer=sqlite_checkpointer)
|
||||
"""
|
||||
pass
|
||||
pass
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timezone, UTC
|
||||
import asyncio
|
||||
|
||||
from src.models import select_model
|
||||
@ -57,5 +57,5 @@ async def agent_cli(agent: BaseAgent, config: RunnableConfig = None):
|
||||
print(f"Tool: {msg.content}")
|
||||
|
||||
def get_cur_time_with_utc():
|
||||
return datetime.now(tz=timezone.utc).isoformat()
|
||||
return datetime.now(tz=UTC).isoformat()
|
||||
|
||||
|
||||
@ -50,7 +50,7 @@ class Config(SimpleConfig):
|
||||
self.add_item("enable_reranker", default=False, des="是否开启重排序")
|
||||
self.add_item("enable_knowledge_base", default=False, des="是否开启知识库")
|
||||
self.add_item("enable_knowledge_graph", default=False, des="是否开启知识图谱")
|
||||
self.add_item("enable_web_search", default=False, des="是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）")
|
||||
self.add_item("enable_web_search", default=False, des="是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）") # noqa: E501
|
||||
# 默认智能体配置
|
||||
self.add_item("default_agent_id", default="", des="默认智能体ID")
|
||||
# 模型配置
|
||||
@ -60,7 +60,7 @@ class Config(SimpleConfig):
|
||||
self.add_item("model_name", default="Qwen/Qwen2.5-7B-Instruct", des="模型名称")
|
||||
|
||||
self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys()))
|
||||
self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys()))
|
||||
self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys())) # noqa: E501
|
||||
self.add_item("model_local_paths", default={}, des="本地模型路径")
|
||||
self.add_item("use_rewrite_query", default="on", des="重写查询", choices=["off", "on", "hyde"])
|
||||
self.add_item("device", default="cuda", des="运行本地模型的设备", choices=["cpu", "cuda"])
|
||||
@ -92,12 +92,12 @@ class Config(SimpleConfig):
|
||||
从 models.yaml 和 models.private.yml 中更新 MODEL_NAMES
|
||||
"""
|
||||
|
||||
with open(Path("src/static/models.yaml"), 'r', encoding='utf-8') as f:
|
||||
with open(Path("src/static/models.yaml"), encoding='utf-8') as f:
|
||||
_models = yaml.safe_load(f)
|
||||
|
||||
# 尝试打开一个 models.private.yml 文件，用来覆盖 models.yaml 中的配置
|
||||
try:
|
||||
with open(Path("src/static/models.private.yml"), 'r', encoding='utf-8') as f:
|
||||
with open(Path("src/static/models.private.yml"), encoding='utf-8') as f:
|
||||
_models_private = yaml.safe_load(f)
|
||||
except FileNotFoundError:
|
||||
_models_private = {}
|
||||
@ -129,7 +129,8 @@ class Config(SimpleConfig):
|
||||
if os.path.exists(self.model_dir):
|
||||
logger.debug(f"MODEL_DIR （{self.model_dir}） 下面的文件夹: {os.listdir(self.model_dir)}")
|
||||
else:
|
||||
logger.warning(f"提醒：MODEL_DIR （{self.model_dir}） 不存在，如果未配置，请忽略，如果配置了，请检查是否配置正确，比如 docker-compose 文件中的映射")
|
||||
logger.warning(f"提醒：MODEL_DIR （{self.model_dir}） 不存在，如果未配置，请忽略，如果配置了，请检查是否配置正确;"
|
||||
"比如 docker-compose 文件中的映射")
|
||||
|
||||
|
||||
# 检查模型提供商是否存在
|
||||
@ -177,7 +178,7 @@ class Config(SimpleConfig):
|
||||
if self.filename is not None and os.path.exists(self.filename):
|
||||
|
||||
if self.filename.endswith(".json"):
|
||||
with open(self.filename, 'r') as f:
|
||||
with open(self.filename) as f:
|
||||
content = f.read()
|
||||
if content:
|
||||
local_config = json.loads(content)
|
||||
@ -186,7 +187,7 @@ class Config(SimpleConfig):
|
||||
print(f"{self.filename} is empty.")
|
||||
|
||||
elif self.filename.endswith(".yaml"):
|
||||
with open(self.filename, 'r') as f:
|
||||
with open(self.filename) as f:
|
||||
content = f.read()
|
||||
if content:
|
||||
local_config = yaml.safe_load(content)
|
||||
|
||||
@ -1,3 +1,3 @@
|
||||
from .history import *
|
||||
from .history import HistoryManager
|
||||
from .knowledgebase import KnowledgeBase
|
||||
from .graphbase import GraphDatabase
|
||||
|
||||
@ -189,7 +189,7 @@ class GraphDatabase:
|
||||
# 筛选出没有embedding的节点
|
||||
nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities)
|
||||
if not nodes_without_embedding:
|
||||
logger.info(f"所有实体已有embedding，无需重新计算")
|
||||
logger.info("所有实体已有embedding，无需重新计算")
|
||||
return
|
||||
|
||||
logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities)}个实体计算embedding")
|
||||
@ -200,7 +200,11 @@ class GraphDatabase:
|
||||
|
||||
for i in range(0, total_entities, max_batch_size):
|
||||
batch_entities = nodes_without_embedding[i:i+max_batch_size]
|
||||
logger.debug(f"Processing entities batch {i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} ({len(batch_entities)} entities)")
|
||||
logger.debug(
|
||||
f"Processing entities batch "
|
||||
f"{i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} "
|
||||
f"({len(batch_entities)} entities)"
|
||||
)
|
||||
|
||||
# 批量获取嵌入向量
|
||||
batch_embeddings = await self.aget_embedding(batch_entities)
|
||||
@ -221,7 +225,7 @@ class GraphDatabase:
|
||||
logger.info(f"Start adding entity to {kgdb_name} with {file_path}")
|
||||
|
||||
def read_triples(file_path):
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
with open(file_path, encoding='utf-8') as file:
|
||||
for line in file:
|
||||
if line.strip():
|
||||
yield json.loads(line.strip())
|
||||
@ -478,7 +482,7 @@ class GraphDatabase:
|
||||
try:
|
||||
graph_info = self.get_graph_info(graph_name)
|
||||
if graph_info is None:
|
||||
logger.error(f"图数据库信息为空，无法保存")
|
||||
logger.error("图数据库信息为空，无法保存")
|
||||
return False
|
||||
|
||||
info_file_path = os.path.join(self.work_dir, "graph_info.json")
|
||||
@ -521,7 +525,7 @@ class GraphDatabase:
|
||||
logger.debug(f"图数据库信息文件不存在：{info_file_path}")
|
||||
return False
|
||||
|
||||
with open(info_file_path, 'r', encoding='utf-8') as f:
|
||||
with open(info_file_path, encoding='utf-8') as f:
|
||||
graph_info = json.load(f)
|
||||
|
||||
# 更新对象属性
|
||||
|
||||
@ -1,7 +1,6 @@
|
||||
from src.utils.prompts import get_system_prompt
|
||||
from src.utils.logging_config import logger
|
||||
|
||||
class HistoryManager():
|
||||
class HistoryManager:
|
||||
def __init__(self, history=None, system_prompt=None):
|
||||
self.messages = []
|
||||
|
||||
@ -47,4 +46,4 @@ class HistoryManager():
|
||||
for message in self.messages:
|
||||
msg = message["content"].replace('\n', ' ')
|
||||
history_str += f"\n{message['role']}: {msg}"
|
||||
return history_str
|
||||
return history_str
|
||||
|
||||
@ -73,11 +73,11 @@ class KnowledgeBase:
|
||||
return
|
||||
|
||||
from src.models.embedding import get_embedding_model
|
||||
self.embed_model = get_embedding_model(config)
|
||||
self.embed_model = get_embedding_model()
|
||||
|
||||
if config.enable_reranker:
|
||||
from src.models.rerank_model import get_reranker
|
||||
self.reranker = get_reranker(config)
|
||||
self.reranker = get_reranker()
|
||||
|
||||
if not self.connect_to_milvus():
|
||||
raise ConnectionError("Failed to connect to Milvus")
|
||||
@ -472,10 +472,7 @@ class KnowledgeBase:
|
||||
async def add_chunks(self, db_id, file_chunks):
|
||||
"""添加分块"""
|
||||
db = self.get_kb_by_id(db_id)
|
||||
|
||||
if db["embed_model"] != self.embed_model.embed_model_fullname:
|
||||
logger.error(f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}")
|
||||
return {"message": f"Embed model not match, cur: {self.embed_model.embed_model_fullname}, req: {db['embed_model']}", "status": "failed"}
|
||||
assert self.check_embed_model(db_id), f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}"
|
||||
|
||||
for file_id, chunk_info in file_chunks.items():
|
||||
# 在数据库中创建文件记录
|
||||
@ -505,10 +502,7 @@ class KnowledgeBase:
|
||||
|
||||
async def add_files(self, db_id, files, params=None):
|
||||
db = self.get_kb_by_id(db_id)
|
||||
|
||||
if not self.check_embed_model(db_id):
|
||||
logger.error(f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}")
|
||||
return {"message": f"Embed model not match, cur: {self.embed_model.embed_model_fullname}, req: {db['embed_model']}", "status": "failed"}
|
||||
assert self.check_embed_model(db_id), f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}"
|
||||
|
||||
# Preprocessing the files to the queue
|
||||
new_files = await self.file_to_chunk(files, params=params)
|
||||
@ -625,11 +619,11 @@ class KnowledgeBase:
|
||||
"embed_model": db["embed_model"],
|
||||
}
|
||||
else:
|
||||
logger.warning((
|
||||
logger.warning(
|
||||
f"无法将知识库 {db['name']} 转换为 Tools, 因为向量模型不匹配，"
|
||||
f"当前向量模型: {self.embed_model.embed_model_fullname}，"
|
||||
f"知识库向量模型: {db['embed_model']}。"
|
||||
))
|
||||
)
|
||||
return retrievers
|
||||
|
||||
################################
|
||||
@ -648,7 +642,8 @@ class KnowledgeBase:
|
||||
logger.info(f"Successfully connected to Milvus at {uri}")
|
||||
return True
|
||||
except MilvusException as e:
|
||||
logger.error(f"Failed to connect to Milvus: {e}，请检查 milvus 的容器是否正常运行，如果已退出，请重新启动 `docker restart milvus-standalone-dev`")
|
||||
logger.error(f"Failed to connect to Milvus: {e}，请检查 milvus 的容器是否正常运行。")
|
||||
logger.error("如果已退出，请重新启动 `docker restart milvus-standalone-dev`。")
|
||||
return False
|
||||
|
||||
def get_collection_names(self):
|
||||
@ -781,4 +776,4 @@ def gen_filename_from_url(url):
|
||||
if len(filename) > 100:
|
||||
filename = filename[:97] + "..."
|
||||
filename = filename.replace('/', '_')
|
||||
return filename
|
||||
return filename
|
||||
|
||||
@ -21,7 +21,7 @@ def migrate_json_to_sqlite():
|
||||
|
||||
try:
|
||||
# 读取旧的JSON文件
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
with open(json_path, encoding='utf-8') as f:
|
||||
old_data = json.load(f)
|
||||
|
||||
if not old_data or not isinstance(old_data, dict) or not old_data.get("databases"):
|
||||
@ -36,7 +36,7 @@ def migrate_json_to_sqlite():
|
||||
logger.info(f"正在迁移知识库: {db_info.get('name')} (ID: {db_id})")
|
||||
|
||||
# 创建知识库
|
||||
db_dict = db_manager.create_database(
|
||||
db_manager.create_database(
|
||||
db_id=db_id,
|
||||
name=db_info.get('name', '未命名知识库'),
|
||||
description=db_info.get('description', ''),
|
||||
@ -81,4 +81,4 @@ def migrate_json_to_sqlite():
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"迁移过程中出错: {e}")
|
||||
return False
|
||||
return False
|
||||
|
||||
@ -13,7 +13,7 @@ class Retriever:
|
||||
|
||||
def _load_models(self):
|
||||
if config.enable_reranker:
|
||||
self.reranker = get_reranker(config)
|
||||
self.reranker = get_reranker()
|
||||
|
||||
if config.enable_web_search:
|
||||
from src.utils.web_search import WebSearcher
|
||||
@ -168,7 +168,7 @@ class Retriever:
|
||||
entities = []
|
||||
if refs["meta"].get("use_graph"):
|
||||
from src.utils.prompts import entity_extraction_prompt_template as entity_template
|
||||
from src.utils.prompts import keywords_prompt_template as entity_template
|
||||
# from src.utils.prompts import keywords_prompt_template as entity_templat|e
|
||||
|
||||
entity_extraction_prompt = entity_template.format(text=query)
|
||||
entities = model.predict(entity_extraction_prompt).content.split("<->")
|
||||
|
||||
@ -4,7 +4,7 @@ from openai import OpenAI
|
||||
from src.utils import logger, get_docker_safe_url
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
class OpenAIBase():
|
||||
class OpenAIBase:
|
||||
def __init__(self, api_key, base_url, model_name, chat_open_ai=None, **kwargs):
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
@ -61,49 +61,6 @@ class OpenAIBase():
|
||||
logger.error(f"Error getting models: {e}")
|
||||
return []
|
||||
|
||||
def _get_model_by_model_url(self, model_url):
|
||||
"""
|
||||
Refs: https://docs.together.ai/reference/models-1
|
||||
|
||||
Return: [
|
||||
{
|
||||
"id": "meta-llama/Meta-Llama-3-70B-Instruct-Turbo",
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"type": "chat",
|
||||
"running": false,
|
||||
"display_name": "Meta Llama 3 70B Instruct Turbo",
|
||||
"organization": "Meta",
|
||||
"link": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct",
|
||||
"license": "Llama-3 (Other)",
|
||||
"context_length": 8192,
|
||||
"config": {
|
||||
"chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}",
|
||||
"stop": [
|
||||
"<|eot_id|>"
|
||||
],
|
||||
"bos_token": "<|begin_of_text|>",
|
||||
"eos_token": "<|end_of_text|>"
|
||||
},
|
||||
"pricing": {
|
||||
"hourly": 0,
|
||||
"input": 0.88,
|
||||
"output": 0.88,
|
||||
"base": 0,
|
||||
"finetune": 0
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
"""
|
||||
headers = {
|
||||
"accept": "application/json",
|
||||
"authorization": f"Bearer {self.api_key}"
|
||||
}
|
||||
response = requests.get(model_url, headers=headers)
|
||||
return response.json()
|
||||
|
||||
|
||||
|
||||
class OpenModel(OpenAIBase):
|
||||
def __init__(self, model_name=None):
|
||||
@ -214,4 +171,4 @@ class DashScope(OpenAIBase):
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
pass
|
||||
|
||||
@ -2,7 +2,6 @@ import os
|
||||
import json
|
||||
import requests
|
||||
import asyncio
|
||||
import json
|
||||
from abc import abstractmethod
|
||||
from zhipuai import ZhipuAI
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
@ -68,7 +67,7 @@ class BaseEmbeddingModel:
|
||||
return data
|
||||
|
||||
class LocalEmbeddingModel(BaseEmbeddingModel):
|
||||
def __init__(self, config, **kwargs):
|
||||
def __init__(self, **kwargs):
|
||||
info = config.embed_model_names[config.embed_model]
|
||||
|
||||
self.model = config.model_local_paths.get(info["name"], info.get("local_path"))
|
||||
@ -104,7 +103,7 @@ class LocalEmbeddingModel(BaseEmbeddingModel):
|
||||
return await self.model.aembed_documents(message)
|
||||
|
||||
def encode_queries(self, queries):
|
||||
logger.warning(f"Huggingface Model 不支持批量 encode queries，因此使用训练实现")
|
||||
logger.warning("Huggingface Model 不支持批量 encode queries，因此使用训练实现")
|
||||
data = []
|
||||
for q in queries:
|
||||
data.append(self.predict(q))
|
||||
@ -114,7 +113,7 @@ class LocalEmbeddingModel(BaseEmbeddingModel):
|
||||
|
||||
class ZhipuEmbedding(BaseEmbeddingModel):
|
||||
|
||||
def __init__(self, config) -> None:
|
||||
def __init__(self) -> None:
|
||||
self.config = config
|
||||
self.model = config.embed_model_names[config.embed_model]["name"]
|
||||
self.dimension = config.embed_model_names[config.embed_model]["dimension"]
|
||||
@ -131,7 +130,7 @@ class ZhipuEmbedding(BaseEmbeddingModel):
|
||||
|
||||
|
||||
class OllamaEmbedding(BaseEmbeddingModel):
|
||||
def __init__(self, config) -> None:
|
||||
def __init__(self) -> None:
|
||||
self.info = config.embed_model_names[config.embed_model]
|
||||
self.model = self.info["name"]
|
||||
self.url = self.info.get("url", "http://localhost:11434/api/embed")
|
||||
@ -155,7 +154,7 @@ class OllamaEmbedding(BaseEmbeddingModel):
|
||||
|
||||
class OtherEmbedding(BaseEmbeddingModel):
|
||||
|
||||
def __init__(self, config) -> None:
|
||||
def __init__(self) -> None:
|
||||
self.info = config.embed_model_names[config.embed_model]
|
||||
self.embed_model_fullname = config.embed_model
|
||||
self.dimension = self.info.get("dimension", None)
|
||||
@ -182,27 +181,28 @@ class OtherEmbedding(BaseEmbeddingModel):
|
||||
"input": message,
|
||||
}
|
||||
|
||||
def get_embedding_model(config):
|
||||
def get_embedding_model():
|
||||
if not config.enable_knowledge_base:
|
||||
return None
|
||||
|
||||
provider, model_name = config.embed_model.split('/', 1)
|
||||
assert config.embed_model in config.embed_model_names.keys(), f"Unsupported embed model: {config.embed_model}, only support {config.embed_model_names.keys()}"
|
||||
support_embed_models = config.embed_model_names.keys()
|
||||
assert config.embed_model in support_embed_models, f"Unsupported embed model: {config.embed_model}, only support {support_embed_models}"
|
||||
logger.debug(f"Loading embedding model {config.embed_model}")
|
||||
if provider == "local":
|
||||
model = LocalEmbeddingModel(config)
|
||||
model = LocalEmbeddingModel()
|
||||
|
||||
elif provider == "zhipu":
|
||||
model = ZhipuEmbedding(config)
|
||||
model = ZhipuEmbedding()
|
||||
|
||||
elif provider == "ollama":
|
||||
model = OllamaEmbedding(config)
|
||||
model = OllamaEmbedding()
|
||||
|
||||
else:
|
||||
model = OtherEmbedding(config)
|
||||
model = OtherEmbedding()
|
||||
|
||||
return model
|
||||
|
||||
def handle_local_model(paths, model_name, default_path):
|
||||
model_path = paths.get(model_name, default_path)
|
||||
return model_path
|
||||
return model_path
|
||||
|
||||
@ -9,7 +9,7 @@ from src.utils.logging_config import logger
|
||||
|
||||
|
||||
class LocalReranker(FlagReranker):
|
||||
def __init__(self, config, **kwargs):
|
||||
def __init__(self, **kwargs):
|
||||
model_info = config.reranker_names[config.reranker]
|
||||
model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("local_path"))
|
||||
model_name_or_path = model_name_or_path or model_info["name"]
|
||||
@ -22,8 +22,8 @@ class LocalReranker(FlagReranker):
|
||||
def sigmoid(x):
|
||||
return 1 / (1 + np.exp(-x))
|
||||
|
||||
class SiliconFlowReranker():
|
||||
def __init__(self, config, **kwargs):
|
||||
class SiliconFlowReranker:
|
||||
def __init__(self, **kwargs):
|
||||
self.url = "https://api.siliconflow.cn/v1/rerank"
|
||||
self.model = config.reranker_names[config.reranker]["name"]
|
||||
|
||||
@ -58,13 +58,14 @@ class SiliconFlowReranker():
|
||||
"max_chunks_per_doc": max_length,
|
||||
}
|
||||
|
||||
def get_reranker(config):
|
||||
assert config.reranker in config.reranker_names.keys(), f"Unsupported Reranker: {config.reranker}, only support {config.reranker_names.keys()}"
|
||||
def get_reranker():
|
||||
support_rerankers = config.reranker_names.keys()
|
||||
assert config.reranker in support_rerankers, f"Unsupported Reranker: {config.reranker}, only support {support_rerankers}"
|
||||
provider, model_name = config.reranker.split('/', 1)
|
||||
if provider == "local":
|
||||
return LocalReranker(config)
|
||||
return LocalReranker()
|
||||
elif provider == "siliconflow":
|
||||
return SiliconFlowReranker(config)
|
||||
return SiliconFlowReranker()
|
||||
else:
|
||||
raise ValueError(f"Unsupported Reranker: {config.reranker}, only support {config.reranker_names.keys()}")
|
||||
|
||||
|
||||
@ -2,4 +2,4 @@ from src.plugins._ocr import OCRPlugin
|
||||
|
||||
ocr = OCRPlugin()
|
||||
|
||||
__all__ = ["ocr"]
|
||||
__all__ = ["ocr"]
|
||||
|
||||
@ -24,7 +24,7 @@ class OCRPlugin:
|
||||
|
||||
def load_model(self):
|
||||
"""加载 OCR 模型"""
|
||||
logger.info(f"加载 OCR 模型，仅在第一次调用时加载")
|
||||
logger.info("加载 OCR 模型，仅在第一次调用时加载")
|
||||
model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR")
|
||||
det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
|
||||
rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")
|
||||
@ -75,7 +75,7 @@ class OCRPlugin:
|
||||
text = '\n'.join([line[1] for line in result])
|
||||
return text
|
||||
else:
|
||||
logger.warning(f"OCR未能识别出文本内容")
|
||||
logger.warning("OCR未能识别出文本内容")
|
||||
return ""
|
||||
|
||||
except Exception as e:
|
||||
@ -214,7 +214,7 @@ def plainreader(file_path):
|
||||
"""读取普通文本文件并返回text文本"""
|
||||
assert os.path.exists(file_path), "File not found"
|
||||
|
||||
with open(file_path, "r") as f:
|
||||
with open(file_path) as f:
|
||||
text = f.read()
|
||||
return text
|
||||
|
||||
|
||||
@ -1,249 +0,0 @@
|
||||
# readmore: https://github.com/zjunlp/DeepKE/blob/main/example/llm/OneKE.md
|
||||
import os
|
||||
import json
|
||||
import torch
|
||||
import dotenv
|
||||
from transformers import (
|
||||
AutoConfig,
|
||||
AutoTokenizer,
|
||||
AutoModelForCausalLM,
|
||||
GenerationConfig,
|
||||
BitsAndBytesConfig
|
||||
)
|
||||
|
||||
from src.utils import logger
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
instruction_mapper = {
|
||||
'NERzh': "你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。",
|
||||
'REzh': "你是专门进行关系抽取的专家。请从input中抽取出符合schema定义的关系三元组。请按照JSON字符串的格式回答。",
|
||||
'EEzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件，不存在的事件返回空列表，不存在的论元返回NAN，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。",
|
||||
'EETzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件类型及事件触发词，不存在的事件返回空列表。请按照JSON字符串的格式回答。",
|
||||
'EEAzh': "你是专门进行事件论元提取的专家。请从input中抽取出符合schema定义的事件论元及论元角色，不存在的论元返回NAN或空字典，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。",
|
||||
'KGzh': '你是一个图谱实体知识结构化专家。根据输入实体类型(entity type)的schema描述，从文本中抽取出相应的实体实例和其属性信息，不存在的属性不输出, 属性存在多值就返回列表，并输出为可解析的json格式。',
|
||||
'NERen': "You are an expert in named entity recognition. Please extract entities that match the schema definition from the input. Return an empty list if the entity type does not exist. Please respond in the format of a JSON string.",
|
||||
'REen': "You are an expert in relationship extraction. Please extract relationship triples that match the schema definition from the input. Return an empty list for relationships that do not exist. Please respond in the format of a JSON string.",
|
||||
'EEen': "You are an expert in event extraction. Please extract events from the input that conform to the schema definition. Return an empty list for events that do not exist, and return NAN for arguments that do not exist. If an argument has multiple values, please return a list. Respond in the format of a JSON string.",
|
||||
'EETen': "You are an expert in event extraction. Please extract event types and event trigger words from the input that conform to the schema definition. Return an empty list for non-existent events. Please respond in the format of a JSON string.",
|
||||
'EEAen': "You are an expert in event argument extraction. Please extract event arguments and their roles from the input that conform to the schema definition, which already includes event trigger words. If an argument does not exist, return NAN or an empty dictionary. Please respond in the format of a JSON string.",
|
||||
'KGen': 'You are an expert in structured knowledge systems for graph entities. Based on the schema description of the input entity type, you extract the corresponding entity instances and their attribute information from the text. Attributes that do not exist should not be output. If an attribute has multiple values, a list should be returned. The results should be output in a parsable JSON format.',
|
||||
}
|
||||
|
||||
split_num_mapper = {
|
||||
'NER':6, 'RE':4, 'EE':4, 'EET':4, 'EEA':4, 'KG':1
|
||||
}
|
||||
|
||||
class OneKE:
|
||||
|
||||
def __init__(self, config=None):
|
||||
|
||||
self.config = config
|
||||
model_name_or_path = config.model_local_paths.get('zjunlp/OneKE', "zjunlp/OneKE")
|
||||
logger.info(f"Loading KGC model OneKE from {model_name_or_path}")
|
||||
|
||||
model_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
|
||||
|
||||
self.quantization_config=BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
llm_int8_threshold=6.0,
|
||||
llm_int8_has_fp16_weight=False,
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
)
|
||||
|
||||
self.generate_config = GenerationConfig(
|
||||
max_length=1024,
|
||||
max_new_tokens=512,
|
||||
return_dict_in_generate=True
|
||||
)
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name_or_path,
|
||||
config=model_config,
|
||||
device_map="auto",
|
||||
quantization_config=self.quantization_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
self.model.eval()
|
||||
|
||||
def construct_input(self, text, schema, task, language="zh", use_split=False):
|
||||
if use_split:
|
||||
split_num = split_num_mapper[task]
|
||||
if isinstance(schema, dict):
|
||||
schema_keys = list(schema.keys())
|
||||
schema_keys = [schema_keys[i:i+split_num] for i in range(0, len(schema_keys), split_num)]
|
||||
schema = [{k: schema[k] for k in keys} for keys in schema_keys]
|
||||
else:
|
||||
schema = [schema[i:i+split_num] for i in range(0, len(schema), split_num)]
|
||||
else:
|
||||
schema = [schema]
|
||||
|
||||
sintructs = []
|
||||
for s in schema:
|
||||
sintruct = json.dumps({
|
||||
"instruction": instruction_mapper[task+language],
|
||||
"schema": s,
|
||||
"input": text,
|
||||
}, ensure_ascii=False)
|
||||
sintructs.append(sintruct)
|
||||
|
||||
return sintructs
|
||||
|
||||
def predict(self, text, schema, task, language="zh", use_split=False):
|
||||
sintructs = self.construct_input(text, schema, task, language, use_split)
|
||||
outputs = []
|
||||
for sintruct in sintructs:
|
||||
input_ids = self.tokenizer.encode(sintruct, return_tensors="pt").to(self.model.device)
|
||||
input_length = input_ids.size(1)
|
||||
|
||||
generation_output = self.model.generate(
|
||||
input_ids=input_ids,
|
||||
generation_config=self.generate_config,
|
||||
pad_token_id=self.tokenizer.eos_token_id
|
||||
)
|
||||
generation_output = generation_output.sequences[0]
|
||||
generation_output = generation_output[input_length:]
|
||||
output = self.tokenizer.decode(generation_output, skip_special_tokens=True)
|
||||
|
||||
outputs.append(output)
|
||||
|
||||
return outputs
|
||||
|
||||
def processing_text_to_kg(self, text_or_path, output_path):
|
||||
for chunk in read_and_process_chars(text_or_path):
|
||||
|
||||
text = chunk
|
||||
schema = [
|
||||
{
|
||||
"entity_type": "食品",
|
||||
"attributes": {
|
||||
"名称": "食品的名称，包括品牌名、通用名称或专业化学名",
|
||||
"分类": "食品所属的类型，例如水果、蔬菜、肉类、谷物、调料、添加剂、益生菌等",
|
||||
"成分": "食品的主要成分，详细列出包括天然成分、添加剂、保鲜剂、营养强化剂等",
|
||||
"营养价值": "食品的营养成分，概括其提供的能量和主要营养素，如蛋白质、脂肪、碳水化合物、维生素和矿物质",
|
||||
"加工方式": "食品的处理或制备方法，包括日常烹饪、加工处理及实验室制备方式等",
|
||||
"作用或食用效果": "食品对健康或身体的影响，可能的功效或用途"
|
||||
}
|
||||
}
|
||||
]
|
||||
task = "KG"
|
||||
output = self.predict(text=text, schema=schema, task=task, language="zh")
|
||||
formatted_output = parse_and_format_output(output=output, task_type=task)
|
||||
|
||||
with open(output_path, 'a+', encoding='utf-8') as f:
|
||||
for entry in formatted_output:
|
||||
f.write(json.dumps(entry, ensure_ascii=False) + '\n')
|
||||
|
||||
print(f"预测结果已添加到 {output_path} 文件中。")
|
||||
return output_path
|
||||
|
||||
def read_and_process_chars(file_path, char_size=512, overlap_size=100):
|
||||
buffer = ""
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
while True:
|
||||
chunk = file.read(char_size)
|
||||
if not chunk: # 文件读取完毕
|
||||
if buffer:
|
||||
yield buffer
|
||||
break
|
||||
chunk = chunk.replace('\n', '').replace('\r', '') # 去除换行符
|
||||
buffer += chunk
|
||||
while len(buffer) >= char_size:
|
||||
yield buffer[:char_size]
|
||||
buffer = buffer[char_size - overlap_size:]
|
||||
|
||||
def parse_and_format_output(output, task_type):
|
||||
formatted_output = []
|
||||
|
||||
for entry in output:
|
||||
try:
|
||||
# Check if the entry is a valid JSON string
|
||||
if isinstance(entry, str) and entry.strip().startswith('{') and entry.strip().endswith('}'):
|
||||
parsed_entry = json.loads(entry)
|
||||
|
||||
if task_type == "KG":
|
||||
for entity_type, entities in parsed_entry.items():
|
||||
for entity_name, attributes in entities.items():
|
||||
for attribute_name, attribute_values in attributes.items():
|
||||
if isinstance(attribute_values, list):
|
||||
for value in attribute_values:
|
||||
formatted_output.append({
|
||||
"h": entity_name,
|
||||
"t": value,
|
||||
"r": attribute_name
|
||||
})
|
||||
else:
|
||||
formatted_output.append({
|
||||
"h": entity_name,
|
||||
"t": attribute_values,
|
||||
"r": attribute_name
|
||||
})
|
||||
elif task_type == "RE":
|
||||
for relation_type, pairs in parsed_entry.items():
|
||||
for pair in pairs:
|
||||
formatted_output.append({
|
||||
"h": pair["subject"],
|
||||
"t": pair["object"],
|
||||
"r": relation_type
|
||||
})
|
||||
else:
|
||||
raise json.JSONDecodeError("Invalid JSON format", entry, 0)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"JSONDecodeError: {e} - Skipping entry")
|
||||
continue
|
||||
except TypeError as e:
|
||||
print(f"TypeError: {e} - Skipping entry")
|
||||
continue
|
||||
except AttributeError as e:
|
||||
print(f"AttributeError: {e} - Skipping entry")
|
||||
continue
|
||||
|
||||
return formatted_output
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
oneke = OneKE()
|
||||
oneke.processing_text_to_kg("asdasdadadsad", 'kg.jsonl')
|
||||
|
||||
file_path = ''
|
||||
output_path = ''
|
||||
text = ""
|
||||
task = "KG"
|
||||
for chunk in read_and_process_chars(file_path):
|
||||
|
||||
text = chunk
|
||||
|
||||
# schema = {
|
||||
# "定义": "描述食品的起源、传统制作方法、文化象征意义或者描述食品或相关事物的定义或含义。包括食品的来源、特点及其在特定文化或背景下的意义。",
|
||||
# "组成成分": "描述食品的组成部分或成分，包括主要成分、微量成分、添加剂等，揭示食品的化学或物理组成。",
|
||||
# "功能": "描述食品或其成分的功能或作用，包括其对人体健康的影响、在烹饪中的用途、药用价值等。",
|
||||
# "属性": "描述食品或其成分的特性或属性，揭示其独特的营养价值、口感特点、保存方式等，包括食品的营养价值、口感特点（如口感丰富、清淡等）、保存方式（如冷藏、冷冻、干燥等）",
|
||||
# "种类": "描述食品或相关事物的种类或类别，揭示其分类体系、不同类型的特点及其在具体应用中的区别。"
|
||||
# }
|
||||
|
||||
schema = [
|
||||
{
|
||||
"entity_type": "食品",
|
||||
"attributes": {
|
||||
"名称": "食品的名称，包括品牌名、通用名称或专业化学名",
|
||||
"分类": "食品所属的类型，例如水果、蔬菜、肉类、谷物、调料、添加剂、益生菌等",
|
||||
"成分": "食品的主要成分，详细列出包括天然成分、添加剂、保鲜剂、营养强化剂等",
|
||||
"营养价值": "食品的营养成分，概括其提供的能量和主要营养素，如蛋白质、脂肪、碳水化合物、维生素和矿物质",
|
||||
"加工方式": "食品的处理或制备方法，包括日常烹饪、加工处理及实验室制备方式等",
|
||||
"作用或食用效果": "食品对健康或身体的影响，可能的功效或用途"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
output = oneke.predict(text=text, schema=schema, task=task, language="zh")
|
||||
formatted_output = parse_and_format_output(output=output, task_type=task)
|
||||
|
||||
with open(output_path, 'a+', encoding='utf-8') as f:
|
||||
for entry in formatted_output:
|
||||
f.write(json.dumps(entry, ensure_ascii=False) + '\n')
|
||||
|
||||
print(f"预测结果已添加到 {output_path} 文件中。")
|
||||
@ -38,4 +38,4 @@ def get_docker_safe_url(base_url):
|
||||
base_url = base_url.replace("http://localhost", "http://host.docker.internal")
|
||||
base_url = base_url.replace("http://127.0.0.1", "http://host.docker.internal")
|
||||
logger.info(f"Running in docker, using {base_url} as base url")
|
||||
return base_url
|
||||
return base_url
|
||||
|
||||
@ -99,4 +99,4 @@ def migrate_knowledge_db():
|
||||
if __name__ == "__main__":
|
||||
# 当作为脚本运行时执行迁移
|
||||
result = migrate_knowledge_db()
|
||||
print(f"迁移{'成功' if result else '失败'}")
|
||||
print(f"迁移{'成功' if result else '失败'}")
|
||||
|
||||
@ -1,8 +1,7 @@
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import time
|
||||
import json
|
||||
from typing import List
|
||||
|
||||
|
||||
async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict:
|
||||
"""发送单个请求到API"""
|
||||
@ -17,6 +16,7 @@ async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict:
|
||||
try:
|
||||
async with session.post(url, json=payload) as response:
|
||||
result = await response.json()
|
||||
print(f"请求 {request_id} 结果: {result}")
|
||||
end_time = time.time()
|
||||
duration = end_time - start_time
|
||||
print(f"请求 {request_id} 完成时间: {time.strftime('%H:%M:%S', time.localtime(end_time))} (耗时: {duration:.2f}秒)")
|
||||
@ -42,13 +42,13 @@ async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict:
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
async def run_concurrent_test(num_requests: int = 10) -> List[dict]:
|
||||
async def run_concurrent_test(num_requests: int = 10) -> list[dict]:
|
||||
"""运行并发测试"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [make_request(session, i) for i in range(num_requests)]
|
||||
return await asyncio.gather(*tasks)
|
||||
|
||||
def analyze_results(results: List[dict]) -> None:
|
||||
def analyze_results(results: list[dict]) -> None:
|
||||
"""分析并打印测试结果"""
|
||||
total_requests = len(results)
|
||||
successful_requests = sum(1 for r in results if r["success"])
|
||||
@ -62,7 +62,7 @@ def analyze_results(results: List[dict]) -> None:
|
||||
else:
|
||||
avg_time = max_time = min_time = 0
|
||||
|
||||
print(f"\n=== 并发测试结果 ===")
|
||||
print("\n=== 并发测试结果 ===")
|
||||
print(f"总请求数: {total_requests}")
|
||||
print(f"成功请求: {successful_requests}")
|
||||
print(f"失败请求: {failed_requests}")
|
||||
@ -80,6 +80,9 @@ def analyze_results(results: List[dict]) -> None:
|
||||
print("\n=== 请求时间线分析 ===")
|
||||
sorted_results = sorted(results, key=lambda x: x["start_time"])
|
||||
test_start_time = sorted_results[0]["start_time"]
|
||||
test_end_time = sorted_results[-1]["end_time"]
|
||||
print(f"测试开始时间: {time.strftime('%H:%M:%S', time.localtime(test_start_time))}")
|
||||
print(f"测试结束时间: {time.strftime('%H:%M:%S', time.localtime(test_end_time))}")
|
||||
|
||||
print("\n时间线详情:")
|
||||
print("请求ID 开始时间 结束时间 耗时(秒) 重叠请求数")
|
||||
|
||||
Loading…
Reference in New Issue
Block a user