Skip to content

L
LAE
Commit 6a19aee9 by tinywell
Options

初步删减代码，仅保留 rate agent 相关部分

parent 5c16c92c
Showing with 39 additions and 2746 deletions
docker-compose.yml
...@@ -7,7 +7,7 @@ services: ...@@ -7,7 +7,7 @@ services:
environment: environment:
POSTGRES_PASSWORD: 111111 POSTGRES_PASSWORD: 111111
ports: ports:
- "5433:5432" - "5434:5432"
volumes: volumes:
- ./lae_pg_data:/var/lib/postgresql/data - ./lae_pg_data:/var/lib/postgresql/data
- ./init_db/create_db.sql:/docker-entrypoint-initdb.d/create_db.sql - ./init_db/create_db.sql:/docker-entrypoint-initdb.d/create_db.sql
......
requirements.txt
...@@ -9,7 +9,6 @@ psycopg2==2.9.7 ...@@ -9,7 +9,6 @@ psycopg2==2.9.7
pydantic==1.10.12 pydantic==1.10.12
requests==2.31.0 requests==2.31.0
sentence-transformers==2.2.2 sentence-transformers==2.2.2
torch==2.0.1
transformers==4.31.0 transformers==4.31.0
uvicorn==0.23.1 uvicorn==0.23.1
unstructured==0.8.1 unstructured==0.8.1
......
src/config/consts.py
# ============================= # =============================
# 资料存储数据库配置 # 资料存储数据库配置
# ============================= # =============================
VEC_DB_HOST = '192.168.10.189' # VEC_DB_HOST = '192.168.10.189'
VEC_DB_HOST = 'localhost'
VEC_DB_DBNAME = 'lae' VEC_DB_DBNAME = 'lae'
VEC_DB_USER = 'postgres' VEC_DB_USER = 'postgres'
VEC_DB_PASSWORD = '111111' VEC_DB_PASSWORD = '111111'
VEC_DB_PORT = '5433' VEC_DB_PORT = '5434'
# ============================= # =============================
# 聊天相关数据库配置 # 聊天相关数据库配置
# ============================= # =============================
CHAT_DB_HOST = '192.168.10.189' # CHAT_DB_HOST = '192.168.10.189'
CHAT_DB_HOST = "localhost"
CHAT_DB_DBNAME = 'lae' CHAT_DB_DBNAME = 'lae'
CHAT_DB_USER = 'postgres' CHAT_DB_USER = 'postgres'
CHAT_DB_PASSWORD = '111111' CHAT_DB_PASSWORD = '111111'
CHAT_DB_PORT = '5433' CHAT_DB_PORT = '5434'
# ============================= # =============================
# 向量化模型路径配置 # 向量化模型路径配置
# ============================= # =============================
EMBEEDING_MODEL_PATH = 'D:\\work\\py\\LAE\\bge-large-zh-v1.5' # EMBEEDING_MODEL_PATH = 'D:\\work\\py\\LAE\\bge-large-zh-v1.5'
EMBEEDING_MODEL_PATH="BAAI/bge-large-zh-v1.5"
# ============================= # =============================
# 重排序模型路径配置 # 重排序模型路径配置
# ============================= # =============================
RERANK_MODEL_PATH = 'D:\\work\\py\\LAE\\bge-reranker-large' # RERANK_MODEL_PATH = 'D:\\work\\py\\LAE\\bge-reranker-large'
# RERANK_MODEL_PATH = 'BAAI/bge-reranker-large' RERANK_MODEL_PATH = 'BAAI/bge-reranker-large'
# ============================= # =============================
# 模型服务URL配置 # 模型服务URL配置
...@@ -41,7 +44,8 @@ SIMILARITY_THRESHOLD = 0.8 ...@@ -41,7 +44,8 @@ SIMILARITY_THRESHOLD = 0.8
# ============================= # =============================
# FAISS向量库文件存储路径配置 # FAISS向量库文件存储路径配置
# ============================= # =============================
FAISS_STORE_PATH = 'D:\\work\\py\\LAE\\faiss' # FAISS_STORE_PATH = 'D:\\work\\py\\LAE\\faiss'
FAISS_STORE_PATH = './tmp/vecstore'
INDEX_NAME = 'know' INDEX_NAME = 'know'
# ============================= # =============================
......
src/controller/api.py
import sys
sys.path.append('../')
from fastapi import FastAPI, Header from fastapi import FastAPI, Header
from src.pgdb.chat.c_db import UPostgresDB
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from src.controller.response import Response
from src.controller.return_data import ReturnData
from src.pgdb.chat.crud import CRUD
import uvicorn import uvicorn
from src.server.qa import QA from src.server.agent_rate import new_rate_agent
from src.pgdb.knowledge.similarity import VectorStore_FAISS from src.controller.request import GeoAgentRateRequest
from src.llm.ernie_with_sdk import ChatERNIESerLLM from langchain_openai import ChatOpenAI
from qianfan import ChatCompletion
from langchain_core.prompts import PromptTemplate
from src.controller.request import (
RegisterRequest,
LoginRequest,
ChatQaRequest,
ChatDetailRequest,
ChatDeleteRequest,
ChatReQA
)
from src.config.consts import (
CHAT_DB_USER,
CHAT_DB_HOST,
CHAT_DB_PORT,
CHAT_DB_DBNAME,
CHAT_DB_PASSWORD,
EMBEEDING_MODEL_PATH,
FAISS_STORE_PATH,
INDEX_NAME,
VEC_DB_HOST,
VEC_DB_PASSWORD,
VEC_DB_PORT,
VEC_DB_USER,
VEC_DB_DBNAME,
SIMILARITY_SHOW_NUMBER,
prompt1
)
app = FastAPI() app = FastAPI()
# 数据库连接
c_db = UPostgresDB(host=CHAT_DB_HOST, database=CHAT_DB_DBNAME, user=CHAT_DB_USER, password=CHAT_DB_PASSWORD,
port=CHAT_DB_PORT, )
c_db.connect()
# 添加 CORS 中间件
app.add_middleware( app.add_middleware(
CORSMiddleware, CORSMiddleware,
allow_origins=["*"], # 允许所有域访问，也可以指定特定域名 allow_origins=["*"], # 允许所有域访问，也可以指定特定域名
...@@ -50,162 +17,29 @@ app.add_middleware( ...@@ -50,162 +17,29 @@ app.add_middleware(
allow_headers=["*"], # 允许所有HTTP头 allow_headers=["*"], # 允许所有HTTP头
) )
base_llm = ChatOpenAI(
openai_api_key='xxxxxxxxxxxxx',
openai_api_base='http://192.168.10.14:8000/v1',
model_name='Qwen2-7B',
verbose=True
)
@app.post("/lae/user/register", response_model=Response) @app.post('/api/agent/rate')
async def register(request: RegisterRequest): def rate(chat_request: GeoAgentRateRequest, token: str = Header(None)):
print(request)
account = request.account
password = request.password
crud = CRUD(_db=c_db)
if crud.user_exist_account(_account=account):
return ReturnData(40000, "当前用户已注册", {}).get_data()
crud.insert_c_user(account=account, password=password)
data = {
"account": account,
}
return ReturnData(200, "用户注册成功", dict(data)).get_data()
@app.get("/lae/user/login", response_model=Response)
async def login(request: LoginRequest):
print(request)
account = request.account
password = request.password
crud = CRUD(_db=c_db)
user_id = crud.user_exist_account_password(_account=account, _password=password)
if not user_id:
return ReturnData(40000, "用户未注册或密码错误", {}).get_data()
token = user_id + '******'
data = {
"account": account,
"token": token
}
return ReturnData(200, "用户登录成功", dict(data)).get_data()
@app.post("/lae/chat/create", response_model=Response)
async def create(token: str = Header(None)):
print(token)
user_id = token.replace('*', '')
crud = CRUD(_db=c_db)
if not crud.user_exist_id(_user_id=user_id):
return ReturnData(40000, "当前用户暂未注册", {}).get_data()
chat_info = '这是该chat的info'
crud.insert_chat(user_id=user_id, info='这是该chat的info', deleted=0)
data = {
"user_id": user_id,
"chat_info": chat_info,
}
return ReturnData(200, '会话创建成功', dict(data)).get_data()
@app.post("/lae/chat/delete", response_model=Response)
async def delete(request: ChatDeleteRequest, token: str = Header(None)):
print(request)
print(token)
user_id = token.replace('*', '')
chat_id = request.chat_id
crud = CRUD(_db=c_db)
if crud.chat_exist_chatid_userid(_chat_id=chat_id, _user_id=user_id):
return ReturnData(40000, "该会话不存在（用户无法访问非本人创建的对话）", {}).get_data()
chat_info = crud.get_chatinfo_from_chatid(chat_id)
crud.delete_chat(chat_id)
data = {
"user_id": user_id,
"chat_info": chat_info,
}
return ReturnData(200, '会话删除成功', dict(data)).get_data()
@app.post("/lae/chat/qa", response_model=Response)
async def qa(request: ChatQaRequest, token: str = Header(None)):
print(request)
print(token)
user_id = token.replace('*', '')
chat_id = request.chat_id
question = request.question
crud = CRUD(_db=c_db)
if not crud.chat_exist_chatid_userid(_chat_id=chat_id, _user_id=user_id):
return ReturnData(40000, "该会话不存在（用户无法访问非本人创建的对话）", {}).get_data()
vecstore_faiss = VectorStore_FAISS(
embedding_model_name=EMBEEDING_MODEL_PATH,
store_path=FAISS_STORE_PATH,
index_name=INDEX_NAME,
info={"port": VEC_DB_PORT, "host": VEC_DB_HOST, "dbname": VEC_DB_DBNAME, "username": VEC_DB_USER,
"password": VEC_DB_PASSWORD},
show_number=SIMILARITY_SHOW_NUMBER,
reset=False)
base_llm = ChatERNIESerLLM(
chat_completion=ChatCompletion(ak="pT7sV1smp4AeDl0LjyZuHBV9", sk="b3N0ibo1IKTLZlSs7weZc8jdR0oHjyMu"))
my_chat = QA(PromptTemplate(input_variables=["context", "question"], template=prompt1), base_llm,
{"temperature": 0.9}, ['context', 'question'], _db=c_db, _chat_id=chat_id,
_faiss_db=vecstore_faiss)
answer = my_chat.chat(question)
my_chat.update_history()
data = {
"answer": answer
}
return ReturnData(200, '模型问答成功', dict(data)).get_data()
@app.get("/lae/chat/detail", response_model=Response)
async def detail(request: ChatDetailRequest, token: str = Header(None)):
print(request)
print(token)
user_id = token.replace('*', '')
chat_id = request.chat_id
crud = CRUD(_db=c_db)
if not crud.chat_exist_chatid_userid(_chat_id=chat_id, _user_id=user_id):
return ReturnData(40000, "该会话不存在（用户无法访问非本人创建的对话）", {}).get_data()
history = crud.get_history(chat_id)
data = {
"chat_id": chat_id,
"history": history
}
return ReturnData(200, "会话详情获取成功", dict(data)).get_data()
@app.post("/lae/chat/clist", response_model=Response) agent = new_rate_agent(base_llm,verbose=True)
async def clist(token: str = Header(None)): try:
print(token) res = agent.exec(prompt_args={"input": chat_request.query})
user_id = token.replace('*', '') except Exception as e:
crud = CRUD(_db=c_db) print(f"处理请求失败, 错误信息: {str(e)}，请重新提问")
chat_list = crud.get_chat_list_userid(user_id) return {
data = { 'code': 500,
"chat_list": chat_list 'data': str(e)
} }
return ReturnData(200, "会话列表获取成功", dict(data)).get_data() return {
'code': 200,
'data': res
@app.post("/lae/chat/reqa", response_model=Response)
async def reqa(request: ChatReQA, token: str = Header(None)):
print(request)
print(token)
chat_id = request.chat_id
user_id = token.replace('*', '')
crud = CRUD(_db=c_db)
if not crud.chat_exist_chatid_userid(_chat_id=chat_id, _user_id=user_id):
return ReturnData(40000, "该会话不存在（用户无法访问非本人创建的对话）", {}).get_data()
question = crud.get_last_question(_chat_id=chat_id)
vecstore_faiss = VectorStore_FAISS(
embedding_model_name=EMBEEDING_MODEL_PATH,
store_path=FAISS_STORE_PATH,
index_name=INDEX_NAME,
info={"port": VEC_DB_PORT, "host": VEC_DB_HOST, "dbname": VEC_DB_DBNAME, "username": VEC_DB_USER,
"password": VEC_DB_PASSWORD},
show_number=SIMILARITY_SHOW_NUMBER,
reset=False)
base_llm = ChatERNIESerLLM(
chat_completion=ChatCompletion(ak="pT7sV1smp4AeDl0LjyZuHBV9", sk="b3N0ibo1IKTLZlSs7weZc8jdR0oHjyMu"))
my_chat = QA(PromptTemplate(input_variables=["context", "question"], template=prompt1), base_llm,
{"temperature": 0.9}, ['context', 'question'], _db=c_db, _chat_id=chat_id,
_faiss_db=vecstore_faiss)
answer = my_chat.chat(question)
my_chat.update_history()
data = {
"answer": answer
} }
return ReturnData(200, '模型重新问答成功', dict(data)).get_data()
if __name__ == "__main__": if __name__ == "__main__":
uvicorn.run(app, host='localhost', port=8889) uvicorn.run(app, host='0.0.0.0', port=8088)
src/controller/web.py deleted 100644 → 0
import sys
sys.path.append('../')
from fastapi import FastAPI, Header,Query
from fastapi.middleware.cors import CORSMiddleware
from datetime import datetime,timedelta
from src.pgdb.chat.c_db import UPostgresDB
import uvicorn
import json
from src.pgdb.chat.crud import CRUD
from src.pgdb.knowledge.similarity import VectorStore_FAISS
from src.pgdb.knowledge.k_db import PostgresDB
from src.pgdb.knowledge.txt_doc_table import TxtDoc
from langchain_openai import ChatOpenAI
from langchain_core.documents import Document
from src.server.rag_query import RagQuery
from src.server.agent_rate import new_rate_agent
from src.controller.request import (
PhoneLoginRequest,
ChatRequest,
ReGenerateRequest,
GeoAgentRateRequest
)
from src.config.consts import (
CHAT_DB_USER,
CHAT_DB_HOST,
CHAT_DB_PORT,
CHAT_DB_DBNAME,
CHAT_DB_PASSWORD,
EMBEEDING_MODEL_PATH,
FAISS_STORE_PATH,
INDEX_NAME,
VEC_DB_HOST,
VEC_DB_PASSWORD,
VEC_DB_PORT,
VEC_DB_USER,
VEC_DB_DBNAME,
SIMILARITY_SHOW_NUMBER,
)
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # 允许所有域访问，也可以指定特定域名
allow_credentials=True,
allow_methods=["*"], # 允许所有HTTP方法
allow_headers=["*"], # 允许所有HTTP头
)
# c_db = UPostgresDB(host=CHAT_DB_HOST, database=CHAT_DB_DBNAME, user=CHAT_DB_USER, password=CHAT_DB_PASSWORD,
# port=CHAT_DB_PORT, )
# c_db.connect()
# k_db = PostgresDB(host=VEC_DB_HOST, database=VEC_DB_DBNAME, user=VEC_DB_USER, password=VEC_DB_PASSWORD, port=VEC_DB_PORT)
# k_db.connect()
# vecstore_faiss = VectorStore_FAISS(
# embedding_model_name=EMBEEDING_MODEL_PATH,
# store_path=FAISS_STORE_PATH,
# index_name=INDEX_NAME,
# info={"port": VEC_DB_PORT, "host": VEC_DB_HOST, "dbname": VEC_DB_DBNAME, "username": VEC_DB_USER,
# "password": VEC_DB_PASSWORD},
# show_number=SIMILARITY_SHOW_NUMBER,
# reset=False)
base_llm = ChatOpenAI(
openai_api_key='xxxxxxxxxxxxx',
openai_api_base='http://192.168.10.14:8000/v1',
model_name='Qwen2-7B',
verbose=True
)
# rag_query = RagQuery(base_llm=base_llm,_faiss_db=vecstore_faiss,_db=TxtDoc(k_db))
@app.post('/api/login')
def login(phone_request: PhoneLoginRequest):
phone = phone_request.phone
crud = CRUD(_db=c_db)
user = crud.user_exist_account(phone)
if not user:
crud.insert_c_user(phone,"123456")
user = crud.user_exist_account(phone)
userid = user[0]
expire = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d %H:%M:%S')
return{
'code': 200,
'data': {
'accessToken': userid,
'refreshToken': userid,
'accessExpire': expire,
'refreshExpire': expire,
'accessUUID': 'accessUUID',
'refreshUUID': 'refreshUUID',
}
}
@app.get('/api/sessions/chat/')
def get_sessions(token: str = Header(None)):
if not token:
return {
'code': 404,
'data': '验证失败'
}
crud = CRUD(_db=c_db)
chat_list = crud.get_chat_list_userid(token)
chat_list_str = []
for chat in chat_list:
chat_list_str.append(str(chat[0]))
return {
'code': 200,
'data': chat_list_str
}
@app.get('/api/session/{session_id}')
def get_history_by_session_id(session_id:str,token: str = Header(None)):
if not token:
return {
'code': 404,
'data': '验证失败'
}
crud = CRUD(_db=c_db)
history = crud.get_history(session_id)
history_json = []
for h in history:
j ={}
j["TurnID"] = h[0]
j["Question"] = h[1]
j["Answer"] = h[2]
j["IsLast"] = h[3]
j["SimilarDocuments"] = get_similarity_doc(h[4])
history_json.append(j)
history_str = json.dumps(history_json)
return {
'code': 200,
'data': history_str
}
@app.delete('/api/session/{session_id}')
def delete_session_by_session_id(session_id:str,token: str = Header(None)):
if not token:
return {
'code': 404,
'data': '验证失败'
}
crud = CRUD(_db=c_db)
crud.delete_chat(session_id)
return {
'code': 200,
'data': 'success'
}
@app.post('/api/general/chat')
def question(chat_request: ChatRequest, token: str = Header(None)):
if not token:
return {
'code': 404,
'data': '验证失败'
}
session_id = chat_request.sessionID
question = chat_request.question
crud = CRUD(_db=c_db)
history = []
if session_id !="":
history = crud.get_last_history(str(session_id))
prompt = ""
for h in history:
prompt += "Q: {}\nA:{}\n".format(h[0], h[1])
res = rag_query.query(question=question,history=prompt)
answer = res["answer"]
docs = res["docs"]
docs_json = json.loads(docs, strict=False)
print(len(docs_json))
doc_hash = []
for d in docs_json:
if "hash" in d:
doc_hash.append(d["hash"])
if len(doc_hash)>0:
hash_str = ",".join(doc_hash)
else:
hash_str = ""
# answer = "test Answer"
if session_id =="":
session_id = crud.create_chat(token, '\t\t', '0')
crud.insert_turn_qa(session_id, question, answer, 0, 1, hash_str)
else:
last_turn_id = crud.get_last_turn_num(str(session_id))
crud.insert_turn_qa(session_id, question, answer, last_turn_id+1, 1, hash_str)
return {
'code': 200,
'data': {
'question': question,
'answer': answer,
'sessionID': session_id,
'similarity': docs_json
}
}
@app.post('/api/general/regenerate')
def re_generate(chat_request: ReGenerateRequest, token: str = Header(None)):
if not token:
return {
'code': 404,
'data': '验证失败'
}
session_id = chat_request.sessionID
question = chat_request.question
crud = CRUD(_db=c_db)
last_turn_id = crud.get_last_turn_num(str(session_id))
history = crud.get_last_history_before_turn_id(str(session_id),last_turn_id)
prompt = ""
for h in history:
prompt += "Q: {}\nA:{}\n".format(h[0], h[1])
res = rag_query.query(question=question, history=prompt)
answer = res["answer"]
docs = res["docs"]
docs_json = json.loads(docs, strict=False)
doc_hash = []
for d in docs_json:
if "hash" in d:
doc_hash.append(d["hash"])
if len(doc_hash)>0:
hash_str = ",".join(doc_hash)
else:
hash_str = ""
# answer = "reGenerate Answer"
crud.update_turn_last(str(session_id), last_turn_id )
crud.insert_turn_qa(session_id, question, answer, last_turn_id, 1, hash_str)
return {
'code': 200,
'data': {
'question': question,
'answer': answer,
'sessionID': session_id,
'similarity': docs_json
}
}
@app.post('/api/agent/rate')
def rate(chat_request: GeoAgentRateRequest, token: str = Header(None)):
agent = new_rate_agent(base_llm,verbose=True)
try:
res = agent.exec(prompt_args={"input": chat_request.query})
except Exception as e:
print(f"处理请求失败, 错误信息: {str(e)}，请重新提问")
return {
'code': 500,
'data': str(e)
}
return {
'code': 200,
'data': res
}
def get_similarity_doc(similarity_doc_hash: str):
if similarity_doc_hash:
hashs = similarity_doc_hash.split(",")
if not similarity_doc_hash or len(hashs) == 0:
return []
docs = []
txt_doc = TxtDoc(k_db)
for h in hashs:
doc = txt_doc.search(h)
if doc is None:
continue
d = Document(page_content=doc[0],metadata=json.loads(doc[1]))
docs.append(d)
return docs_to_json(docs)
def docs_to_json(docs):
docs_json = []
for d in docs:
j = {}
j["page_content"] = d.page_content
j["from_file"] = d.metadata["filename"]
j["page_number"] = 0
docs_json.append(j)
return docs_json
if __name__ == "__main__":
uvicorn.run(app, host='0.0.0.0', port=8088)
src/gradio_app.py deleted 100644 → 0
# -*- coding: utf-8 -*-
import sys
sys.path.append('../')
import gradio as gr
from langchain_core.prompts import PromptTemplate
from src.llm.chatglm import ChatGLMSerLLM
from src.llm.ernie_with_sdk import ChatERNIESerLLM
from qianfan import ChatCompletion
import os
import asyncio
from src.pgdb.chat.c_db import UPostgresDB
from src.pgdb.knowledge.similarity import VectorStore_FAISS
from src.config.consts import (
CHAT_DB_USER,
CHAT_DB_HOST,
CHAT_DB_PORT,
CHAT_DB_DBNAME,
CHAT_DB_PASSWORD,
EMBEEDING_MODEL_PATH,
FAISS_STORE_PATH,
INDEX_NAME,
VEC_DB_HOST,
VEC_DB_PASSWORD,
VEC_DB_PORT,
VEC_DB_USER,
VEC_DB_DBNAME,
SIMILARITY_SHOW_NUMBER,
GR_PORT,
GR_SERVER_NAME,
ICON_PATH
)
from src.server.qa import QA
prompt1 = """'''
{context}
'''
请你根据上述已知资料回答下面的问题，问题如下：
{question}"""
PROMPT1 = PromptTemplate(input_variables=["context", "question"], template=prompt1)
def main():
c_db = UPostgresDB(host=CHAT_DB_HOST, database=CHAT_DB_DBNAME, user=CHAT_DB_USER, password=CHAT_DB_PASSWORD,
port=CHAT_DB_PORT, )
vecstore_faiss = VectorStore_FAISS(
embedding_model_name=EMBEEDING_MODEL_PATH,
store_path=FAISS_STORE_PATH,
index_name=INDEX_NAME,
info={"port": VEC_DB_PORT, "host": VEC_DB_HOST, "dbname": VEC_DB_DBNAME, "username": VEC_DB_USER,
"password": VEC_DB_PASSWORD},
show_number=SIMILARITY_SHOW_NUMBER,
reset=False)
base_llm = ChatERNIESerLLM(
chat_completion=ChatCompletion(ak="pT7sV1smp4AeDl0LjyZuHBV9", sk="b3N0ibo1IKTLZlSs7weZc8jdR0oHjyMu"))
# base_llm = ChatGLMSerLLM(url='http://192.168.22.106:8088')
my_chat = QA(PROMPT1, base_llm, {"temperature": 0.9}, ['context', 'question'], _db=c_db,
_faiss_db=vecstore_faiss)
def clear(): # 清空输入框
return ''
def show_history(): # 显示对话历史
return my_chat.get_history()
def stop_btn():
return gr.Button(interactive=False)
def restart_btn():
return gr.Button(interactive=True)
def get_users():
o_users = my_chat.get_users()
users_l = [item[0] for item in o_users]
return gr.components.Radio(choices=users_l, label="选择一个用户", value=users_l[0], interactive=True), users_l[
0]
def create_chat(user_account):
my_chat.create_chat(user_account)
def get_chats(user_account):
o_chats = my_chat.get_chats(user_account)
if len(o_chats) >= 13:
o_chats = o_chats[:13]
chats_l = [item[0] + ':' + item[1] for item in o_chats]
if my_chat.chat_id:
result = [item for item in chats_l if item.split(":")[0].strip() == my_chat.chat_id][0]
return gr.components.Radio(choices=chats_l, label="历史对话", value=result, interactive=True,
show_label=True)
else:
return gr.components.Radio(choices=chats_l, label="历史对话", value=chats_l[0], interactive=True,
show_label=True)
def set_info(question):
my_chat.set_info(question)
def set_chat_id(chat_id_info):
chat_id = chat_id_info.split(':')[0]
my_chat.set_chat_id(chat_id)
def load():
l_users, t_user = get_users()
l_chats = get_chats(t_user)
return l_users, l_chats
def clear_text(t):
if t == "请输入您的问题":
return ""
else:
return t
def reset_text():
return ""
def blur(t):
if t == "":
return "请输入您的问题"
else:
return t
def text_his(text, history):
history = history + [[text, None]]
return history
def clear_tip(his):
if not his[0][0] and his[0][1] == "你好，我是新晨科技股份的低空经济人工智能助手小晨，如果您有低空经济相关问题，欢迎随时向我咨询。":
return his[1:]
else:
return his
with gr.Blocks(css='index.css', title="低空经济知识问答") as demo:
gr.HTML("""<h1 align="center">低空经济知识问答</h1>""", visible=False)
with gr.Row():
with gr.Column(scale=2, visible=False):
users = gr.components.Radio(choices=[], label="选择一个用户", interactive=True,
visible=False, show_label=False)
chats = gr.components.Radio(choices=[], label="历史对话", interactive=True,
show_label=True, visible=False)
new_chat_btn = gr.Button("新建对话", visible=False)
with gr.Column(scale=8):
chatbot = gr.Chatbot(bubble_full_width=False,
avatar_images=(ICON_PATH + '\\user2.png', ICON_PATH + "\\bot2.png"),
value=[[None,
"你好，我是新晨科技股份的低空经济人工智能助手小晨，如果您有低空经济相关问题，欢迎随时向我咨询。"]],
height=400, show_copy_button=True,
show_label=False, line_breaks=True)
with gr.Row():
input_text = gr.Textbox(show_label=False, lines=1, label="文本输入", scale=9, container=False,
placeholder="请输入您的问题", max_lines=1)
sub_btn = gr.Button("提交", scale=1)
sub_btn.click(
stop_btn, [], sub_btn
).success(
clear_tip, [chatbot], [chatbot]
).success(
text_his, [input_text, chatbot], [chatbot]
).success(
reset_text, [], input_text
).success(
my_chat.async_chat2, [chatbot], [chatbot]
).success(
restart_btn, [], sub_btn
)
# input_text.submit(
# stop_btn, [], sub_btn
# ).then(
# my_chat.async_chat2, [input_text, chatbot], [chatbot]
# ).then(
# restart_btn, [], sub_btn
# )
demo.load(load, [], [users, chats])
# input_text.submit(my_chat.async_chat, [input_text], [chatbot]
# ).then(
# stop_btn, None, sub_btn
# ).then(
# set_info, [input_text], []
# ).then(
# get_chats, [users], [chats]
# ).then(
# my_chat.update_history, None, None
# ).then(
# show_history, None, chatbot
# ).then(
# clear, None, [input_text]
# ).then(
# restart_btn, None, sub_btn
# ).then(
# reset_text, [], input_text
# )
# new_chat_btn.click(create_chat, [users], []).then(
# get_chats, [users], [chats]
# )
# users.change(get_chats, [users], [chats]).then(
# set_chat_id, [chats], None
# ).then(
# show_history, None, chatbot
# )
# chats.change(set_chat_id, [chats], None).then(
# show_history, None, chatbot
# )
# sub_btn.click(my_chat.async_chat, [input_text], [chatbot]
# ).then(
# stop_btn, None, sub_btn
# ).then(
# set_info, [input_text], []
# ).then(
# get_chats, [users], [chats]
# ).then(
# my_chat.update_history, None, None
# ).then(
# show_history, None, chatbot
# ).then(
# clear, None, [input_text]
# ).then(
# restart_btn, None, sub_btn
# ).then(
# reset_text, [], input_text
# )
demo.queue().launch(share=False, inbrowser=True, server_name='0.0.0.0', server_port=GR_PORT)
if __name__ == "__main__":
main()
src/index.css deleted 100644 → 0
#component-5 {
height: 76vh !important;
overflow: auto !important;
}
.wrap.svelte-vm32wk.svelte-vm32wk.svelte-vm32wk {
display: inline !important;
}
.wrap .svelte-vm32wk label {
margin-bottom: 10px !important;
}
#component-9 {
height: 88vh !important;
border: #f6faff;
box-shadow: none;
background: #f6faff;
}
footer {
visibility: hidden;
}
.app.svelte-1kyws56.svelte-1kyws56 {
height: 100vh;
}
@media screen and (max-width: 768px) {
input[type="text"], textarea {
-webkit-user-modify: read-write-plaintext-only;
-webkit-text-size-adjust: none;
}
#component-3 {
display: none;
}
#component-5 {
display: none !important;
}
#component-9 {
height: 89vh !important;
}
.app.svelte-1kyws56.svelte-1kyws56 {
height: 100vh;
}
#component-10 {
position: fixed;
bottom: 15px;
right: 0px;
padding: 0 20px;
}
#component-12 {
min-width: min(74px, 100%);
}
}
#component-3 button {
background: #4999ff !important;
color: white !important;
}
#component-3 button:hover {
background: #4999ff !important;
color: white !important;
}
span.svelte-1gfkn6j {
margin: 55px 0;
}
gradio-app {
background-color: #f6faff !important;
}
.bot.svelte-1pjfiar.svelte-1pjfiar.svelte-1pjfiar {
background: white;
border: none;
box-shadow: 0px 0px 9px 0px rgba(0, 0, 0, 0.1);
border-radius: 4px;
}
.user.svelte-1pjfiar.svelte-1pjfiar.svelte-1pjfiar {
background: white;
border: none;
box-shadow: 0px 0px 9px 0px rgba(0, 0, 0, 0.1);
border-radius: 4px;
}
.message-buttons-bubble.svelte-1pjfiar.svelte-1pjfiar.svelte-1pjfiar {
/* background: white; */
border: none;
box-shadow: 0px 0px 9px 0px rgba(0, 0, 0, 0.1);
}
label.svelte-vm32wk > .svelte-vm32wk + .svelte-vm32wk {
color: #989898;
}
p {
color: #26415f;
}
.wrapper.svelte-nab2ao {
background: #f6faff;
}
#component-12 {
background: #4999ff;
color: white;
}
.svelte-11hlfrc svg {
color: gray;
}
#component-6 {
position: fixed;
width: 219px;
height: 36px;
border-radius: 18px;
top: 4%;
z-index: 4;
left: 8%;
}
div.svelte-sfqy0y {
background-color: white;
}
src/llm/__init__.py deleted 100644 → 0
"""各种大模型提供的服务"""
src/llm/chatglm.py deleted 100644 → 0
import os
import requests
from typing import Dict, Optional, List, Any, Mapping, Iterator
from pydantic import root_validator
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig
import langchain
from langchain_core.language_models import BaseLLM, LLM
from langchain_community.cache import InMemoryCache
from langchain.callbacks.manager import Callbacks
from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun
import aiohttp
import asyncio
# 启动llm的缓存
# langchain.llm_cache = InMemoryCache()
class ChatGLMLocLLM(LLM):
model_name: str = "THUDM/chatglm-6b"
ptuning_checkpoint: str = None
quantization_bit: Optional[int] = None
pre_seq_len: Optional[int] = None
prefix_projection: bool = False
tokenizer: AutoTokenizer = None
model: AutoModel = None
def _llm_type(self) -> str:
return "chatglm_local"
# @root_validator()
@staticmethod
def validate_environment(values: Dict) -> Dict:
if not values["model_name"]:
raise ValueError("No model name provided.")
model_name = values["model_name"]
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
# model = AutoModel.from_pretrained(model_name, config=config, trust_remote_code=True)
if values["pre_seq_len"]:
config.pre_seq_len = values["pre_seq_len"]
if values["prefix_projection"]:
config.prefix_projection = values["prefix_projection"]
if values["ptuning_checkpoint"]:
ptuning_checkpoint = values["ptuning_checkpoint"]
print(f"Loading prefix_encoder weight from {ptuning_checkpoint}")
model = AutoModel.from_pretrained(model_name, config=config, trust_remote_code=True)
prefix_state_dict = torch.load(os.path.join(ptuning_checkpoint, "pytorch_model.bin"))
new_prefix_state_dict = {}
for k, v in prefix_state_dict.items():
if k.startswith("transformer.prefix_encoder."):
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
else:
model = AutoModel.from_pretrained(model_name, config=config, trust_remote_code=True).half().cuda()
if values["pre_seq_len"]:
# P-tuning v2
model = model.half().cuda()
model.transformer.prefix_encoder.float().cuda()
if values["quantization_bit"]:
print(f"Quantized to {values['quantization_bit']} bit")
model = model.quantize(values["quantization_bit"])
model = model.eval()
values["tokenizer"] = tokenizer
values["model"] = model
return values
def _call(self, prompt: str, stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs) -> str:
resp, his = self.model.chat(self.tokenizer, prompt)
# print(f"prompt:{prompt}\nresponse:{resp}\n")
return resp
class ChatGLMSerLLM(LLM):
# 模型服务url
url: str = "http://127.0.0.1:8000"
chat_history: dict = []
out_stream: bool = False
cache: bool = False
@property
def _llm_type(self) -> str:
return "chatglm3-6b"
def get_num_tokens(self, text: str) -> int:
resp = self._post(url=self.url + "/tokens", query=self._construct_query(text))
if resp.status_code == 200:
resp_json = resp.json()
predictions = resp_json['response']
# display(self.convert_data(resp_json['history']))
return predictions
else:
return len(text)
@staticmethod
def convert_data(data):
result = []
for item in data:
result.append({'q': item[0], 'a': item[1]})
return result
def _construct_query(self, prompt: str, temperature=0.95) -> Dict:
"""构造请求体
"""
# self.chat_history.append({"role": "user", "content": prompt})
query = {
"prompt": prompt,
"history": self.chat_history,
"max_length": 4096,
"top_p": 0.7,
"temperature": temperature
}
return query
@classmethod
def _post(cls, url: str,
query: Dict) -> Any:
"""POST请求
"""
_headers = {"Content_Type": "application/json"}
with requests.session() as sess:
resp = sess.post(url,
json=query,
headers=_headers,
timeout=300)
return resp
@staticmethod
async def _post_stream(url: str,
query: Dict,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, stream=False) -> Any:
"""POST请求
"""
_headers = {"Content_Type": "application/json"}
async with aiohttp.ClientSession() as sess:
async with sess.post(url, json=query, headers=_headers, timeout=300) as response:
if response.status == 200:
if stream and not run_manager:
print('not callable')
if run_manager:
for _callable in run_manager.get_sync().handlers:
await _callable.on_llm_start(None, None)
async for chunk in response.content.iter_any():
# 处理每个块的数据
if chunk and run_manager:
for _callable in run_manager.get_sync().handlers:
# print(chunk.decode("utf-8"),end="")
await _callable.on_llm_new_token(chunk.decode("utf-8"))
if run_manager:
for _callable in run_manager.get_sync().handlers:
await _callable.on_llm_end(None)
else:
raise ValueError(f'glm 请求异常，http code：{response.status}')
def _call(self, prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
stream=False,
**kwargs: Any) -> str:
query = self._construct_query(prompt=prompt,
temperature=kwargs["temperature"] if "temperature" in kwargs else 0.95)
# display("==============================")
# display(query)
# post
if stream or self.out_stream:
async def _post_stream():
await self._post_stream(url=self.url + "/stream",
query=query, run_manager=run_manager, stream=stream or self.out_stream)
asyncio.run(_post_stream())
return ''
else:
resp = self._post(url=self.url,
query=query)
if resp.status_code == 200:
resp_json = resp.json()
# self.chat_history.append({'q': prompt, 'a': resp_json['response']})
predictions = resp_json['response']
# display(self.convert_data(resp_json['history']))
return predictions
else:
raise ValueError(f'glm 请求异常，http code：{resp.status_code}')
async def _acall(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
query = self._construct_query(prompt=prompt,
temperature=kwargs["temperature"] if "temperature" in kwargs else 0.95)
await self._post_stream(url=self.url + "/stream",
query=query, run_manager=run_manager, stream=self.out_stream)
return ''
@property
def _identifying_params(self) -> Mapping[str, Any]:
"""Get the identifying parameters.
"""
_param_dict = {
"url": self.url
}
return _param_dict
src/llm/chatglm_openapi.py deleted 100644 → 0
import os
import requests
from typing import Dict, Optional, List, Any, Mapping, Iterator
from pydantic import root_validator
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig
import langchain
from langchain_core.language_models import BaseLLM, LLM
from langchain_openai import OpenAI
from langchain_community.cache import InMemoryCache
from langchain.callbacks.manager import Callbacks
from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun
class ChatGLMSerLLM(OpenAI):
def get_token_ids(self, text: str) -> List[int]:
if self.model_name.__contains__("chatglm"):
## 发起http请求，获取token_ids
url = f"{self.openai_api_base}/num_tokens"
query = {"prompt": text, "model": self.model_name}
_headers = {"Content_Type": "application/json", "Authorization": "chatglm " + self.openai_api_key}
resp = self._post(url=url, query=query, headers=_headers)
if resp.status_code == 200:
resp_json = resp.json()
print(resp_json)
predictions = resp_json['choices'][0]['text']
## predictions字符串转int
return [int(predictions)]
return [len(text)]
@classmethod
def _post(cls, url: str,
query: Dict, headers: Dict) -> Any:
"""POST请求
"""
_headers = {"Content_Type": "application/json"}
_headers.update(headers)
with requests.session() as sess:
resp = sess.post(url,
json=query,
headers=_headers,
timeout=300)
return resp
src/llm/ernie.py deleted 100644 → 0
import logging
import os
from typing import Any, Dict, List, Mapping, Optional
from langchain_core.language_models import BaseLLM, LLM
from langchain_core.outputs import LLMResult
from langchain_core.utils import get_from_dict_or_env
from langchain.callbacks.manager import Callbacks
from langchain_core.callbacks import CallbackManagerForLLMRun
from enum import Enum
from pydantic import root_validator, Field
from .ernie_sdk import CompletionRequest, ErnieBot, Message, bot_message, user_message
logger = logging.getLogger(__name__)
class ModelType(Enum):
ERNIE = "ernie"
ERNIE_LITE = "ernie-lite"
SHEETS1 = "sheets1"
SHEETS2 = "sheets2"
SHEET_COMB = "sheet-comb"
LLAMA2_7B = "llama2-7b"
LLAMA2_13B = "llama2-13b"
LLAMA2_70B = "llama2-70b"
QFCN_LLAMA2_7B = "qfcn-llama2-7b"
BLOOMZ_7B = "bloomz-7b"
MODEL_SERVICE_BASE_URL = "https://aip.baidubce.com/rpc/2.0/"
MODEL_SERVICE_Suffix = {
ModelType.ERNIE: "ai_custom/v1/wenxinworkshop/chat/completions",
ModelType.ERNIE_LITE: "ai_custom/v1/wenxinworkshop/chat/eb-instant",
ModelType.SHEETS1: "ai_custom/v1/wenxinworkshop/chat/besheet",
ModelType.SHEETS2: "ai_custom/v1/wenxinworkshop/chat/besheets2",
ModelType.SHEET_COMB: "ai_custom/v1/wenxinworkshop/chat/sheet_comb1",
ModelType.LLAMA2_7B: "ai_custom/v1/wenxinworkshop/chat/llama_2_7b",
ModelType.LLAMA2_13B: "ai_custom/v1/wenxinworkshop/chat/llama_2_13b",
ModelType.LLAMA2_70B: "ai_custom/v1/wenxinworkshop/chat/llama_2_70b",
ModelType.QFCN_LLAMA2_7B: "ai_custom/v1/wenxinworkshop/chat/qianfan_chinese_llama_2_7b",
ModelType.BLOOMZ_7B: "ai_custom/v1/wenxinworkshop/chat/bloomz_7b1",
}
class ErnieLLM(LLM):
"""
ErnieLLM is a LLM that uses Ernie to generate text.
"""
model_name: Optional[ModelType] = None
access_token: Optional[str] = ""
@root_validator()
def validate_environment(self, values: Dict) -> Dict:
"""Validate the environment."""
# print(values)
model_name = ModelType(get_from_dict_or_env(values, "model_name", "model_name", str(ModelType.ERNIE)))
access_token = get_from_dict_or_env(values, "access_token", "ERNIE_ACCESS_TOKEN", "")
if not access_token:
raise ValueError("No access token provided.")
values["model_name"] = model_name
values["access_token"] = access_token
return values
def _call(self, prompt: str, stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs) -> str:
request = CompletionRequest(messages=[Message("user", prompt)])
bot = ErnieBot(_get_model_service_url(self.model_name), self.access_token or "", request)
try:
# 你的代码
response = bot.get_response().result
# print("response: ",response)
return response
except Exception as e:
# 处理异常
print("exception:", e)
return e.__str__()
@property
def _llm_type(self) -> str:
"""Return type of llm."""
return "ernie"
# def _identifying_params(self) -> Mapping[str, Any]:
# return {
# "name": "ernie",
# }
def _get_model_service_url(model_name) -> str:
# print("_get_model_service_url model_name: ",model_name)
return MODEL_SERVICE_BASE_URL + MODEL_SERVICE_Suffix[model_name]
class ErnieChat(LLM):
model_name: ModelType
access_token: str
prefix_messages: List = Field(default_factory=list)
id: str = ""
def _call(self, prompt: str, stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs) -> str:
msg = user_message(prompt)
request = CompletionRequest(messages=self.prefix_messages + [msg])
bot = ErnieBot(_get_model_service_url(self.model_name), self.access_token, request)
try:
# 你的代码
response = bot.get_response().result
if self.id == "":
self.id = bot.get_response().id
self.prefix_messages.append(msg)
self.prefix_messages.append(bot_message(response))
return response
except Exception as e:
# 处理异常
raise e
def _get_id(self) -> str:
return self.id
@property
def _llm_type(self) -> str:
"""Return type of llm."""
return "ernie"
src/llm/ernie_sdk.py deleted 100644 → 0
from dataclasses import asdict, dataclass
from typing import List
from pydantic import BaseModel, Field
from enum import Enum
class MessageRole(str, Enum):
USER = "user"
BOT = "assistant"
@dataclass
class Message:
role: str
content: str
@dataclass
class CompletionRequest:
messages: List[Message]
stream: bool = False
user: str = ""
@dataclass
class Usage:
prompt_tokens: int
completion_tokens: int
total_tokens: int
@dataclass
class CompletionResponse:
id: str
object: str
created: int
result: str
need_clear_history: bool
ban_round: int = 0
sentence_id: int = 0
is_end: bool = False
usage: Usage = None
is_safe: bool = False
is_truncated: bool = False
class ErrorResponse(BaseModel):
error_code: int = Field(...)
error_msg: str = Field(...)
id: str = Field(...)
class ErnieBot:
url: str
access_token: str
request: CompletionRequest
def __init__(self, url: str, access_token: str, request: CompletionRequest):
self.url = url
self.access_token = access_token
self.request = request
def get_response(self) -> CompletionResponse:
import requests
import json
headers = {'Content-Type': 'application/json'}
params = {'access_token': self.access_token}
request_dict = asdict(self.request)
response = requests.post(self.url, params=params, data=json.dumps(request_dict), headers=headers)
# print(response.json())
try:
return CompletionResponse(**response.json())
except Exception as e:
print(e)
raise Exception(response.json())
def user_message(prompt: str) -> Message:
return Message(MessageRole.USER, prompt)
def bot_message(prompt: str) -> Message:
return Message(MessageRole.BOT, prompt)
src/llm/ernie_with_sdk.py deleted 100644 → 0
import os
import requests
from typing import Dict, Optional, List, Any, Mapping, Iterator
from pydantic import root_validator
from langchain_core.language_models import LLM
from langchain_community.cache import InMemoryCache
from langchain.callbacks.manager import Callbacks
from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun
import qianfan
from qianfan import ChatCompletion
# 启动llm的缓存
# langchain.llm_cache = InMemoryCache()
class ChatERNIESerLLM(LLM):
# 模型服务url
chat_completion: ChatCompletion = None
# url: str = "http://127.0.0.1:8000"
chat_history: dict = []
out_stream: bool = False
cache: bool = False
model_name: str = "ERNIE-Bot"
# def __init__(self):
# self.chat_completion = qianfan.ChatCompletion(ak="pT7sV1smp4AeDl0LjyZuHBV9", sk="b3N0ibo1IKTLZlSs7weZc8jdR0oHjyMu")
@property
def _llm_type(self) -> str:
return self.model_name
def get_num_tokens(self, text: str) -> int:
return len(text)
@staticmethod
def convert_data(data):
result = []
for item in data:
result.append({'q': item[0], 'a': item[1]})
return result
def _call(self, prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
stream=False,
**kwargs: Any) -> str:
resp = self.chat_completion.do(model=self.model_name, messages=[{
"role": "user",
"content": prompt
}])
print(resp)
assert resp.code == 200
return resp.body["result"]
async def _post_stream(self,
query: Dict,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
stream=False) -> Any:
"""POST请求
"""
async for r in await self.chat_completion.ado(model=self.model_name, messages=[query], stream=stream):
assert r.code == 200
if run_manager:
for _callable in run_manager.get_sync().handlers:
await _callable.on_llm_new_token(r.body["result"])
async def _acall(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
await self._post_stream(query={
"role": "user",
"content": prompt
}, stream=True, run_manager=run_manager)
return ''
src/llm/loader.py deleted 100644 → 0
import os
import transformers
import torch
from transformers import AutoModel, AutoTokenizer, AutoConfig, DataCollatorForSeq2Seq
from peft import PeftModel
class ModelLoader:
def __init__(self, model_name_or_path, pre_seq_len=0, prefix_projection=False):
self.config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
if pre_seq_len is not None and pre_seq_len > 0:
self.config.pre_seq_len = pre_seq_len
self.config.prefix_projection = prefix_projection
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
self.model = AutoModel.from_pretrained(model_name_or_path, config=self.config, trust_remote_code=True).half()
# self.model = self.model.cuda()
self.base_model = self.model
def quantize(self, quantization_bit):
if quantization_bit is not None:
print(f"Quantized to {quantization_bit} bit")
self.model = self.model.quantize(quantization_bit)
return self.model
def models(self):
return self.model, self.tokenizer
def collator(self):
return DataCollatorForSeq2Seq(tokenizer=self.tokenizer, model=self.model)
def load_lora(self, ckpt_path, name="default"):
# 训练时节约GPU占用
_peft_loaded = PeftModel.from_pretrained(self.base_model, ckpt_path, adapter_name=name)
self.model = _peft_loaded.merge_and_unload()
print(f"Load LoRA model successfully!")
def load_loras(self, ckpt_paths, name="default"):
global peft_loaded
if len(ckpt_paths) == 0:
return
first = True
for name, path in ckpt_paths.items():
print(f"Load {name} from {path}")
if first:
peft_loaded = PeftModel.from_pretrained(self.base_model, path, adapter_name=name)
first = False
else:
peft_loaded.load_adapter(path, adapter_name=name)
peft_loaded.set_adapter(name)
self.model = peft_loaded
def load_prefix(self, ckpt_path):
prefix_state_dict = torch.load(os.path.join(ckpt_path, "pytorch_model.bin"))
new_prefix_state_dict = {}
for k, v in prefix_state_dict.items():
if k.startswith("transformer.prefix_encoder."):
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
self.model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
self.model.transformer.prefix_encoder.float()
print(f"Load prefix model successfully!")
src/loader/__init__.py deleted 100644 → 0
"""资料分割"""
\ No newline at end of file
src/loader/callback.py deleted 100644 → 0
from abc import ABC, abstractmethod
class BaseCallback(ABC):
@abstractmethod
def filter(self, title: str, content: str) -> bool: # return True舍弃当前段落
pass
src/loader/chinese_text_splitter.py deleted 100644 → 0
import re
from typing import List
from src.loader.config import SENTENCE_SIZE
from langchain_text_splitters import CharacterTextSplitter
class ChineseTextSplitter(CharacterTextSplitter):
def __init__(self, pdf: bool = False, sentence_size: int = SENTENCE_SIZE, **kwargs):
super().__init__(**kwargs)
self.pdf = pdf
self.sentence_size = sentence_size
def split_text1(self, text: str) -> List[str]:
if self.pdf:
text = re.sub(r"\n{3,}", "\n", text)
text = re.sub('\s', ' ', text)
text = text.replace("\n\n", "")
sent_sep_pattern = re.compile('([﹒﹔﹖﹗．。！？]["’”」』]{0,2}|(?=["‘“「『]{1,2}|$))') # del ：；
sent_list = []
for ele in sent_sep_pattern.split(text):
if sent_sep_pattern.match(ele) and sent_list:
sent_list[-1] += ele
elif ele:
sent_list.append(ele)
return sent_list
def split_text(self, text: str) -> List[str]: ##此处需要进一步优化逻辑
if self.pdf:
text = re.sub(r"\n{3,}", r"\n", text)
text = re.sub('\s', " ", text)
text = re.sub("\n\n", "", text)
text = re.sub(r'([;；.!?。！？\?])([^”’])', r"\1\n\2", text) # 单字符断句符
text = re.sub(r'(\.{6})([^"’”」』])', r"\1\n\2", text) # 英文省略号
text = re.sub(r'(\…{2})([^"’”」』])', r"\1\n\2", text) # 中文省略号
text = re.sub(r'([;；!?。！？\?]["’”」』]{0,2})([^;；!?，。！？\?])', r'\1\n\2', text)
# 如果双引号前有终止符，那么双引号才是句子的终点，把分句符\n放到双引号后，注意前面的几句都小心保留了双引号
text = text.rstrip() # 段尾如果有多余的\n就去掉它
# 很多规则中会考虑分号;，但是这里我把它忽略不计，破折号、英文双引号等同样忽略，需要的再做些简单调整即可。
ls = [i for i in text.split("\n") if i]
for ele in ls:
if len(ele) > self.sentence_size:
ele1 = re.sub(r'([,，.]["’”」』]{0,2})([^,，.])', r'\1\n\2', ele)
ele1_ls = ele1.split("\n")
for ele_ele1 in ele1_ls:
if len(ele_ele1) > self.sentence_size:
ele_ele2 = re.sub(r'([\n]{1,}| {2,}["’”」』]{0,2})([^\s])', r'\1\n\2', ele_ele1)
ele2_ls = ele_ele2.split("\n")
for ele_ele2 in ele2_ls:
if len(ele_ele2) > self.sentence_size:
ele_ele3 = re.sub('( ["’”」』]{0,2})([^ ])', r'\1\n\2', ele_ele2)
ele2_id = ele2_ls.index(ele_ele2)
ele2_ls = ele2_ls[:ele2_id] + [i for i in ele_ele3.split("\n") if i] + ele2_ls[
ele2_id + 1:]
ele_id = ele1_ls.index(ele_ele1)
ele1_ls = ele1_ls[:ele_id] + [i for i in ele2_ls if i] + ele1_ls[ele_id + 1:]
_id = ls.index(ele)
ls = ls[:_id] + [i for i in ele1_ls if i] + ls[_id + 1:]
return ls
src/loader/config.py deleted 100644 → 0
# 文本分句长度
SENTENCE_SIZE = 100
ZH_TITLE_ENHANCE = False
src/loader/zh_title_enhance.py deleted 100644 → 0
from typing import List
from langchain_core.documents import Document
import re
def under_non_alpha_ratio(text: str, threshold: float = 0.5):
"""Checks if the proportion of non-alpha characters in the text snippet exceeds a given
threshold. This helps prevent text like "-----------BREAK---------" from being tagged
as a title or narrative text. The ratio does not count spaces.
Parameters
----------
text
The input string to test
threshold
If the proportion of non-alpha characters exceeds this threshold, the function
returns False
"""
if len(text) == 0:
return False
alpha_count = len([char for char in text if char.strip() and char.isalpha()])
total_count = len([char for char in text if char.strip()])
try:
ratio = alpha_count / total_count
return ratio < threshold
except:
return False
def is_possible_title(
text: str,
title_max_word_length: int = 20,
non_alpha_threshold: float = 0.5,
) -> bool:
"""Checks to see if the text passes all the checks for a valid title.
Parameters
----------
text
The input text to check
title_max_word_length
The maximum number of words a title can contain
non_alpha_threshold
The minimum number of alpha characters the text needs to be considered a title
"""
# 文本长度为0的话，肯定不是title
if len(text) == 0:
print("Not a title. Text is empty.")
return False
# 文本中有标点符号，就不是title
ENDS_IN_PUNCT_PATTERN = r"[^\w\s]\Z"
ENDS_IN_PUNCT_RE = re.compile(ENDS_IN_PUNCT_PATTERN)
if ENDS_IN_PUNCT_RE.search(text) is not None:
return False
# 文本长度不能超过设定值，默认20
# NOTE(robinson) - splitting on spaces here instead of word tokenizing because it
# is less expensive and actual tokenization doesn't add much value for the length check
if len(text) > title_max_word_length:
return False
# 文本中数字的占比不能太高，否则不是title
if under_non_alpha_ratio(text, threshold=non_alpha_threshold):
return False
# NOTE(robinson) - Prevent flagging salutations like "To My Dearest Friends," as titles
if text.endswith((",", ".", "，", "。")):
return False
if text.isnumeric():
print(f"Not a title. Text is all numeric:\n\n{text}") # type: ignore
return False
# 开头的字符内应该有数字，默认5个字符内
if len(text) < 5:
text_5 = text
else:
text_5 = text[:5]
alpha_in_text_5 = sum(list(map(lambda x: x.isnumeric(), list(text_5))))
if not alpha_in_text_5:
return False
return True
def zh_title_enhance(docs: List[Document]):
title = None
if len(docs) > 0:
for doc in docs:
if is_possible_title(doc.page_content):
doc.metadata['category'] = 'cn_Title'
title = doc.page_content
elif title:
doc.page_content = f"下文与({title})有关。{doc.page_content}"
return docs
else:
print("文件不存在")
src/model_api/run_chatglm2.sh deleted 100644 → 0
# shellcheck disable=SC1128
#!/bin/bash
# Set the path to the server.py script
SERVER_PATH=server.py
# Set the default values for the arguments
MODEL_NAME_OR_PATH="../../../model/chatglm2-6b"
CHECKPOINT=None
CHECKPOINT_PATH="../../../model/ckpt/chatglm2-6b-qlora-INSv11-rank16-1e-3-30/checkpoint-2000"
PRE_SEQ_LEN=128
QUANTIZATION_BIT=8
PORT=8002
# Call the server.py script with the parsed arguments
python $SERVER_PATH \
--model_name_or_path $MODEL_NAME_OR_PATH \
--checkpoint $CHECKPOINT \
--checkpoint_path $CHECKPOINT_PATH \
--pre_seq_len $PRE_SEQ_LEN \
--quantization_bit $QUANTIZATION_BIT \
--port $PORT
\ No newline at end of file
src/model_api/run_spdsvb.sh deleted 100644 → 0
#!/bin/bash
# Set the path to the server.py script
SERVER_PATH=server.py
# Set the default values for the arguments
MODEL_NAME_OR_PATH="../../../model/chatglm2-6b"
CHECKPOINT=lora
CHECKPOINT_PATH="../../../model/ckpt/chatglm2-6b-qlora-INSv11-rank16-1e-3-30/checkpoint-2000"
QUANTIZATION_BIT=8
PORT=8001
# Call the server.py script with the parsed arguments
python $SERVER_PATH \
--model_name_or_path $MODEL_NAME_OR_PATH \
--checkpoint $CHECKPOINT \
--checkpoint_path $CHECKPOINT_PATH \
--quantization_bit $QUANTIZATION_BIT \
--port $PORT
\ No newline at end of file
src/model_api/server.py deleted 100644 → 0
import argparse
import time
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import json
import datetime
import torch
from typing import AsyncIterable
from pydantic import BaseModel
import uvicorn
import signal
from src.llm.loader import ModelLoader
DEVICE = "cuda"
DEVICE_ID = "0"
CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE
def torch_gc():
if torch.cuda.is_available():
with torch.cuda.device(CUDA_DEVICE):
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def build_history(history):
result = []
for item in history if history else []:
result.append((item['q'], item['a']))
return result
def convert_data(data):
result = []
for item in data:
result.append({'q': item[0], 'a': item[1]})
return result
class StreamRequest(BaseModel):
"""Request body for streaming."""
message: str
stop_stream = False
def signal_handler(signal, frame):
global stop_stream
stop_stream = True
async def send_message(message: str, history=[], max_length=2048, top_p=0.7, temperature=0.95) -> AsyncIterable[str]:
global model, tokenizer, stop_stream
count = 0
old_len = 0
print(message)
output = ''
for response, history in model.stream_chat(tokenizer, message, history=history,
max_length=max_length,
top_p=top_p,
temperature=temperature):
# print(old_len,count)
if stop_stream:
stop_stream = False
break
else:
output = response[old_len:]
print(output, end='', flush=True)
# print(output)
old_len = len(response)
signal.signal(signal.SIGINT, signal_handler)
yield f"{output}"
print("")
# yield f"\n"
# print()
app = FastAPI()
@app.post("/stream")
async def stream(request: Request):
json_post_raw = await request.json()
json_post = json.dumps(json_post_raw)
json_post_list = json.loads(json_post)
prompt = json_post_list.get('prompt')
history = build_history(json_post_list.get('history'))
max_length = json_post_list.get('max_length')
top_p = json_post_list.get('top_p')
temperature = json_post_list.get('temperature')
return StreamingResponse(send_message(prompt, history=history, max_length=max_length if max_length else 2048,
top_p=top_p if top_p else 0.7,
temperature=temperature if temperature else 0.95), media_type="text/plain")
@app.post("/")
async def create_item(request: Request):
global model, tokenizer
json_post_raw = await request.json()
json_post = json.dumps(json_post_raw)
json_post_list = json.loads(json_post)
prompt = json_post_list.get('prompt')
history = build_history(json_post_list.get('history'))
max_length = json_post_list.get('max_length')
top_p = json_post_list.get('top_p')
temperature = json_post_list.get('temperature')
response, history = model.chat(tokenizer,
prompt,
history=history,
max_length=max_length if max_length else 2048,
top_p=top_p if top_p else 0.7,
temperature=temperature if temperature else 0.95)
now = datetime.datetime.now()
time = now.strftime("%Y-%m-%d %H:%M:%S")
answer = {
"response": response,
"history": history,
"status": 200,
"time": time
}
log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
print(log)
torch_gc()
return answer
@app.post("/tokens")
async def get_num_tokens(request: Request):
global model, tokenizer
json_post_raw = await request.json()
json_post = json.dumps(json_post_raw)
json_post_list = json.loads(json_post)
prompt = json_post_list.get('prompt')
tokens = tokenizer.encode(prompt, add_special_tokens=False)
print("=======================================")
print("=======================================")
print(len(tokens), prompt)
print("=======================================")
print("=======================================")
now = datetime.datetime.now()
time = now.strftime("%Y-%m-%d %H:%M:%S")
answer = {
"response": len(tokens),
"status": 200,
"time": time
}
return answer
def parse_args():
parser = argparse.ArgumentParser(description='ChatGLM2-6B Server')
parser.add_argument('--model_name_or_path', type=str, default='THUDM/chatglm2-6b', help='模型id或local path')
parser.add_argument('--checkpoint', type=str, default=None, help='checkpoint类型（None、ptuning、lora）')
parser.add_argument('--checkpoint_path', type=str,
default='../../../model/ckpt/chatglm2-6b-qlora-INSv11-rank16-1e-3-30/checkpoint-2000',
help='checkpoint路径')
parser.add_argument('--pre_seq_len', type=int, default=128, help='prefix 长度')
parser.add_argument('--quantization_bit', type=int, default=None, help='是否量化')
parser.add_argument('--port', type=int, default=8000, help='端口')
parser.add_argument('--host', type=str, default='0.0.0.0', help='host')
# parser.add_argument('--max_input_length', type=int, default=512, help='instruction + input的最大长度')
# parser.add_argument('--max_output_length', type=int, default=1536, help='output的最大长度')
return parser.parse_args()
if __name__ == '__main__':
cfg = parse_args()
## ----------- load model --------------
start = time.time()
if cfg.checkpoint == "lora":
# lora 微调 checkpoint 及模型加载
loader = ModelLoader(cfg.model_name_or_path)
loader.load_lora(cfg.checkpoint_path)
elif cfg.checkpoint == "ptuning":
# ptuning v2 微调 checkpoint 及模型加载
loader = ModelLoader(cfg.model_name_or_path, cfg.pre_seq_len, False)
loader.load_prefix(cfg.checkpoint_path)
else:
loader = ModelLoader(cfg.model_name_or_path)
model, tokenizer = loader.models()
if cfg.quantization_bit is not None:
model = loader.quantize(cfg.quantization_bit)
model.cuda().eval()
uvicorn.run(app, host=cfg.host, port=cfg.port, workers=1)
src/pgdb/chat/__init__.py deleted 100644 → 0
"""会话信息相关表"""
\ No newline at end of file
src/pgdb/chat/c_db.py deleted 100644 → 0
import psycopg2
from psycopg2 import OperationalError, InterfaceError
class UPostgresDB:
"""
psycopg2.connect(
dsn #指定连接参数。可以使用参数形式或 DSN 形式指定。
host #指定连接数据库的主机名。
dbname #指定数据库名。
user #指定连接数据库使用的用户名。
password #指定连接数据库使用的密码。
port #指定连接数据库的端口号。
connection_factory #指定创建连接对象的工厂类。
cursor_factory #指定创建游标对象的工厂类。
async_ #指定是否异步连接(默认False)。
sslmode #指定 SSL 模式。
sslrootcert #指定证书文件名。
sslkey #指定私钥文件名。
sslcert #指定公钥文件名。
)
"""
def __init__(self, host, database, user, password, port=5432):
self.host = host
self.database = database
self.user = user
self.password = password
self.port = port
self.conn = None
self.cur = None
def connect(self):
try:
self.conn = psycopg2.connect(
host=self.host,
database=self.database,
user=self.user,
password=self.password,
port=self.port
)
self.cur = self.conn.cursor()
except Exception as e:
print(f"连接数据库出现错误: {e}")
def execute(self, query):
try:
if self.conn is None or self.conn.closed:
self.connect()
self.cur.execute(query)
self.conn.commit()
except InterfaceError as e:
print(f"数据库连接已经关闭: {e}")
except OperationalError as e:
print(f"数据库连接出现问题: {e}")
self.connect()
self.retry_execute(query)
except Exception as e:
print(f"执行sql语句出现错误: {e}")
self.conn.rollback()
def retry_execute(self, query):
try:
self.cur.execute(query)
self.conn.commit()
except Exception as e:
print(f"重新执行sql语句再次出现错误: {type(e).__name__}: {e}")
self.conn.rollback()
def execute_args(self, query, args):
try:
if self.conn is None or self.conn.closed:
self.connect()
self.cur.execute(query, args)
self.conn.commit()
except InterfaceError as e:
print(f"数据库连接已经关闭: {e}")
except OperationalError as e:
print(f"数据库操作出现问题: {e}")
self.connect()
self.retry_execute_args(query, args)
except Exception as e:
print(f"执行sql语句出现错误: {e}")
self.conn.rollback()
def retry_execute_args(self, query, args):
try:
self.cur.execute(query, args)
self.conn.commit()
except Exception as e:
print(f"重新执行sql语句再次出现错误: {type(e).__name__}: {e}")
self.conn.rollback()
def search(self, query, params=None):
if self.conn is None or self.conn.closed:
self.connect()
self.cur.execute(query, params)
def fetchall(self):
return self.cur.fetchall()
def fetchone(self):
return self.cur.fetchone()
def close(self):
self.cur.close()
self.conn.close()
def format(self, query):
try:
if self.conn is None or self.conn.closed:
self.connect()
self.cur.execute(query)
self.conn.commit()
except Exception as e:
print(f"An error occurred: {e}")
self.conn.rollback()
src/pgdb/chat/c_user_table.py deleted 100644 → 0
from .c_db import UPostgresDB
TABLE_USER = """
DROP TABLE IF EXISTS "c_user";
CREATE TABLE c_user (
user_id varchar(1000) PRIMARY KEY,
account varchar(20) NOT NULL,
password varchar(50) NOT NULL
);
COMMENT ON COLUMN "c_user"."user_id" IS '用户id';
COMMENT ON COLUMN "c_user"."account" IS '用户帐户';
COMMENT ON COLUMN "c_user"."password" IS '用户密码';
COMMENT ON TABLE "c_user" IS '用户表';
"""
class CUser:
def __init__(self, _user_id, _account, _password) -> None:
self.user_id = _user_id
self.account = _account
self.password = _password
src/pgdb/chat/chat_table.py deleted 100644 → 0
from .c_db import UPostgresDB
TABLE_CHAT = """
DROP TABLE IF EXISTS "chat";
CREATE TABLE chat (
chat_id varchar(1000) PRIMARY KEY,
user_id int,
info text,
create_time timestamp(6) DEFAULT current_timestamp,
deleted int2
);
COMMENT ON COLUMN "chat"."chat_id" IS '会话id';
COMMENT ON COLUMN "chat"."user_id" IS '会话创建用户id';
COMMENT ON COLUMN "chat"."info" IS '会话简介';
COMMENT ON COLUMN "chat"."create_time" IS '会话创建时间，默认为当前时间';
COMMENT ON COLUMN "chat"."deleted" IS '是否删除：0=否，1=是';
COMMENT ON TABLE "chat" IS '会话信息表';
"""
class Chat:
def __init__(self, _chat_id, _user_id, _info, _create_time, _deleted) -> None:
self.chat_id = _chat_id
self.user_id = _user_id
self.info = _info
self.create_time = _create_time
self.deleted = _deleted
src/pgdb/chat/crud.py deleted 100644 → 0
from .c_db import UPostgresDB
TABLE_CHAT = """
DROP TABLE IF EXISTS "chat";
CREATE TABLE chat (
chat_id varchar(1000) PRIMARY KEY,
user_id varchar(1000),
info text,
create_time timestamp(6) DEFAULT current_timestamp,
deleted int2
);
COMMENT ON COLUMN "chat"."chat_id" IS '会话id';
COMMENT ON COLUMN "chat"."user_id" IS '会话创建用户id';
COMMENT ON COLUMN "chat"."info" IS '会话简介';
COMMENT ON COLUMN "chat"."create_time" IS '会话创建时间，默认为当前时间';
COMMENT ON COLUMN "chat"."deleted" IS '是否删除：0=否，1=是';
COMMENT ON TABLE "chat" IS '会话信息表';
DROP SEQUENCE IF EXISTS "chat_seq";
CREATE SEQUENCE "chat_seq"
INCREMENT 1
MINVALUE 1
MAXVALUE 2147483647
START 1
CACHE 1;
ALTER TABLE chat ALTER COLUMN chat_id SET DEFAULT nextval('chat_seq'::regclass);
"""
TABLE_TURN_QA = """
DROP TABLE IF EXISTS "turn_qa";
CREATE TABLE turn_qa (
turn_id varchar(1000) PRIMARY KEY,
chat_id varchar(1000),
question text,
answer text,
similar_docs text,
create_time timestamp(6) DEFAULT current_timestamp,
turn_number int,
is_last int2
);
COMMENT ON COLUMN "turn_qa"."turn_id" IS '会话轮次id';
COMMENT ON COLUMN "turn_qa"."chat_id" IS '会话id';
COMMENT ON COLUMN "turn_qa"."question" IS '该轮会话问题';
COMMENT ON COLUMN "turn_qa"."answer" IS '该轮会话答案';
COMMENT ON COLUMN "turn_qa"."similar_docs" IS '该轮会话相似文档 hash 索引';
COMMENT ON COLUMN "turn_qa"."create_time" IS '该轮会话创建时间，默认为当前时间';
COMMENT ON COLUMN "turn_qa"."turn_number" IS '会话轮数';
COMMENT ON COLUMN "turn_qa"."is_last" IS '是否为最后一轮对话：0=否，1=是';
COMMENT ON TABLE "turn_qa" IS '会话轮次信息表';
DROP SEQUENCE IF EXISTS "turn_qa_seq";
CREATE SEQUENCE "turn_qa_seq"
INCREMENT 1
MINVALUE 1
MAXVALUE 2147483647
START 1
CACHE 1;
ALTER TABLE turn_qa ALTER COLUMN turn_id SET DEFAULT nextval('turn_qa_seq'::regclass);
"""
TABLE_USER = """
DROP TABLE IF EXISTS "c_user";
CREATE TABLE c_user (
user_id varchar(1000) PRIMARY KEY,
account varchar(20) NOT NULL,
password varchar(50) NOT NULL
);
COMMENT ON COLUMN "c_user"."user_id" IS '用户id';
COMMENT ON COLUMN "c_user"."account" IS '用户帐户';
COMMENT ON COLUMN "c_user"."password" IS '用户密码';
COMMENT ON TABLE "c_user" IS '用户表';
DROP SEQUENCE IF EXISTS "c_user_seq";
CREATE SEQUENCE "c_user_seq"
INCREMENT 1
MINVALUE 1
MAXVALUE 2147483647
START 1
CACHE 1;
ALTER TABLE c_user ALTER COLUMN user_id SET DEFAULT nextval('c_user_seq'::regclass);
"""
class CRUD:
def __init__(self, _db: UPostgresDB):
self.db = _db
def create_table(self):
self.db.execute(TABLE_CHAT)
self.db.execute(TABLE_TURN_QA)
self.db.execute(TABLE_USER)
def get_history(self, _chat_id):
query = f'SELECT turn_number,question,answer,is_last,similar_docs FROM turn_qa WHERE chat_id=(%s) ORDER BY turn_number ASC'
self.db.execute_args(query, (_chat_id,))
ans = self.db.fetchall()
return ans
def get_last_history(self, _chat_id):
query = f'SELECT question,answer,similar_docs FROM turn_qa WHERE chat_id=(%s) and is_last=1 ORDER BY turn_number ASC'
self.db.execute_args(query, (_chat_id,))
ans = self.db.fetchall()
return ans
def get_last_history_before_turn_id(self, _chat_id,turn_id):
query = f'SELECT question,answer,similar_docs FROM turn_qa WHERE chat_id=(%s) and is_last=1 and turn_number<(%s) ORDER BY turn_number ASC'
self.db.execute_args(query, (_chat_id,turn_id))
ans = self.db.fetchall()
return ans
def insert_turn_qa(self, chat_id, question, answer, turn_number, is_last, similar_docs=None):
query = f'INSERT INTO turn_qa(chat_id, question, answer, turn_number, is_last, similar_docs) VALUES (%s,%s,%s,%s,%s,%s)'
self.db.execute_args(query, (chat_id, question, answer, turn_number, is_last, similar_docs))
def insert_c_user(self, account, password):
query = f'INSERT INTO c_user(account, password) VALUES (%s,%s)'
self.db.execute_args(query, (account, password))
def insert_chat(self, user_id, info, deleted):
query = f'INSERT INTO chat(user_id, info, deleted) VALUES (%s,%s,%s)'
self.db.execute_args(query, (user_id, info, deleted))
def update_last(self, chat_id):
query = f'UPDATE turn_qa SET is_last = 0 WHERE chat_id = (%s) AND is_last = 1'
self.db.execute_args(query, (chat_id,))
def update_turn_last(self, chat_id, turn_number):
query = f'UPDATE turn_qa SET is_last = 0 WHERE chat_id = (%s) AND turn_number = (%s)'
self.db.execute_args(query, (chat_id, turn_number))
def user_exist_id(self, _user_id):
query = f'SELECT * FROM c_user WHERE user_id = (%s)'
self.db.execute_args(query, (_user_id,))
return self.db.fetchone()
def user_exist_account(self, _account):
query = f'SELECT * FROM c_user WHERE account = (%s)'
self.db.execute_args(query, (_account,))
return self.db.fetchone()
def user_exist_account_password(self, _account, _password):
query = f'SELECT user_id FROM c_user WHERE account = (%s) AND password = (%s)'
self.db.execute_args(query, (_account, _password))
return self.db.fetchone()
def chat_exist_chatid_userid(self, _chat_id, _user_id):
query = f'SELECT * FROM chat WHERE chat_id = (%s) AND user_id = (%s)'
self.db.execute_args(query, (_chat_id, _user_id))
return self.db.fetchone()
def get_chat_list_userid(self, _user_id):
query = f'SELECT chat_id FROM chat WHERE user_id = (%s) AND deleted = 0 order by create_time desc'
self.db.execute_args(query, (_user_id,))
return self.db.fetchall()
def get_chatinfo_from_chatid(self, _chat_id):
query = f'SELECT info FROM chat WHERE chat_id = (%s)'
self.db.execute_args(query, (_chat_id,))
return self.db.fetchone()
def delete_chat(self, _chat_id):
query = f'UPDATE chat SET deleted = 1 WHERE chat_id = (%s)'
self.db.execute_args(query, (_chat_id,))
def get_last_question(self, _chat_id):
query = f'SELECT question FROM turn_qa WHERE chat_id = (%s) AND turn_number = 1'
self.db.execute_args(query, (_chat_id,))
return self.db.fetchone()[0]
def get_users(self):
query = f'SELECT account FROM c_user'
self.db.execute(query)
return self.db.fetchall()
def get_chats(self, account):
query = f'SELECT chat.chat_id,chat.info FROM chat JOIN c_user ON chat.user_id = c_user.user_id WHERE c_user.account = (%s) ORDER BY chat.create_time DESC;'
self.db.execute_args(query, (account,))
return self.db.fetchall()
def create_chat(self, user_id, info, deleted):
query = f'INSERT INTO chat(user_id, info, deleted) VALUES (%s,%s,%s) RETURNING chat_id'
self.db.execute_args(query, (user_id, info, deleted))
ans = self.db.fetchall()[0][0]
return ans
def get_uersid_from_account(self, account):
query = f'SELECT user_id FROM c_user WHERE account = (%s)'
self.db.execute_args(query, (account, ))
ans = self.db.fetchall()[0][0]
print(ans)
return ans
def get_chat_info(self, chat_id):
query = f'SELECT info FROM chat WHERE chat_id = (%s)'
self.db.execute_args(query, (chat_id,))
ans = self.db.fetchall()[0][0]
print(ans)
return ans
def set_info(self, chat_id, info):
query = f'UPDATE chat SET info = (%s) WHERE chat_id = (%s)'
self.db.execute_args(query, (info, chat_id))
def get_last_turn_num(self,chat_id):
query = f'SELECT max(turn_number) FROM turn_qa WHERE chat_id = (%s)'
self.db.execute_args(query, (chat_id,))
ans = self.db.fetchone()[0]
return ans
src/pgdb/chat/turn_qa_table.py deleted 100644 → 0
TABLE_TURN_QA = """
DROP TABLE IF EXISTS "turn_qa";
CREATE TABLE turn_qa (
turn_id varchar(1000) PRIMARY KEY,
chat_id varchar(1000),
question text,
answer text,
similar_docs text,
create_time timestamp(6) DEFAULT current_timestamp,
turn_number int,
is_last int2
);
COMMENT ON COLUMN "turn_qa"."turn_id" IS '会话轮次id';
COMMENT ON COLUMN "turn_qa"."chat_id" IS '会话id';
COMMENT ON COLUMN "turn_qa"."question" IS '该轮会话问题';
COMMENT ON COLUMN "turn_qa"."answer" IS '该轮会话答案';
COMMENT ON COLUMN "turn_qa"."similar_docs" IS '该轮会话相似文档 hash 索引';
COMMENT ON COLUMN "turn_qa"."create_time" IS '该轮会话创建时间，默认为当前时间';
COMMENT ON COLUMN "turn_qa"."turn_number" IS '会话轮数';
COMMENT ON COLUMN "turn_qa"."is_last" IS '是否为最后一轮对话：0=否，1=是';
COMMENT ON TABLE "turn_qa" IS '会话轮次信息表';
"""
class TurnQa:
def __init__(self, _turn_id, _chat_id, _question, _answer, _create_time, _turn_number, _is_last) -> None:
self.turn_id = _turn_id
self.chat_id = _chat_id
self.question = _question
self.answer = _answer
self.create_time = _create_time
self.turn_number = _turn_number
self.is_last = _is_last
\ No newline at end of file
src/pgdb/knowledge/__init__.py deleted 100644 → 0
"""资料存储相关"""
\ No newline at end of file
src/pgdb/knowledge/callback.py deleted 100644 → 0
import sys
from abc import ABC, abstractmethod
import json
from typing import List, Tuple
from langchain_core.documents import Document
from src.pgdb.knowledge.pgsqldocstore import PgSqlDocstore, str2hash_base64
sys.path.append("../")
class DocumentCallback(ABC):
@abstractmethod # 向量库储存前文档处理--
def before_store(self, docstore: PgSqlDocstore, documents):
pass
@abstractmethod # 向量库查询后文档处理--用于结构建立
def after_search(self, docstore: PgSqlDocstore, documents: List[Tuple[Document, float]], number: int = 1000) -> \
List[Tuple[Document, float]]: # 向量库查询后文档处理
pass
class DefaultDocumentCallback(DocumentCallback):
def before_store(self, docstore: PgSqlDocstore, documents):
output_doc = []
for doc in documents:
if "next_doc" in doc.metadata:
doc.metadata["next_hash"] = str2hash_base64(doc.metadata["next_doc"])
doc.metadata.pop("next_doc")
output_doc.append(doc)
return output_doc
def after_search(self, docstore: PgSqlDocstore, documents: List[Tuple[Document, float]], number: int = 1000) -> \
List[Tuple[Document, float]]: # 向量库查询后文档处理
output_doc: List[Tuple[Document, float]] = []
exist_hash = []
for doc, score in documents:
print(exist_hash)
dochash = str2hash_base64(doc.page_content)
if dochash in exist_hash:
continue
else:
exist_hash.append(dochash)
output_doc.append((doc, score))
if len(output_doc) > number:
return output_doc
fordoc = doc
while "next_hash" in fordoc.metadata:
if len(fordoc.metadata["next_hash"]) > 0:
if fordoc.metadata["next_hash"] in exist_hash:
break
else:
exist_hash.append(fordoc.metadata["next_hash"])
content = docstore.TXT_DOC.search(fordoc.metadata["next_hash"])
if content:
fordoc = Document(page_content=content[0], metadata=json.loads(content[1]))
output_doc.append((fordoc, score))
if len(output_doc) > number:
return output_doc
else:
break
else:
break
return output_doc
src/pgdb/knowledge/k_db.py deleted 100644 → 0
import psycopg2
class PostgresDB:
"""
psycopg2.connect(
dsn #指定连接参数。可以使用参数形式或 DSN 形式指定。
host #指定连接数据库的主机名。
dbname #指定数据库名。
user #指定连接数据库使用的用户名。
password #指定连接数据库使用的密码。
port #指定连接数据库的端口号。
connection_factory #指定创建连接对象的工厂类。
cursor_factory #指定创建游标对象的工厂类。
async_ #指定是否异步连接(默认False)。
sslmode #指定 SSL 模式。
sslrootcert #指定证书文件名。
sslkey #指定私钥文件名。
sslcert #指定公钥文件名。
)
"""
def __init__(self, host, database, user, password, port=5432):
self.host = host
self.database = database
self.user = user
self.password = password
self.port = port
self.conn = None
self.cur = None
def connect(self):
self.conn = psycopg2.connect(
host=self.host,
database=self.database,
user=self.user,
password=self.password,
port=self.port
)
self.cur = self.conn.cursor()
def execute(self, query):
try:
self.cur.execute(query)
self.conn.commit()
except Exception as e:
print(f"An error occurred: {e}")
self.conn.rollback()
def execute_args(self, query, args):
try:
self.cur.execute(query, args)
self.conn.commit()
except Exception as e:
print(f"An error occurred: {e}")
self.conn.rollback()
def search(self, query, params=None):
self.cur.execute(query, params)
def fetchall(self):
return self.cur.fetchall()
def close(self):
self.cur.close()
self.conn.close()
def format(self, query):
try:
self.cur.execute(query)
self.conn.commit()
except Exception as e:
print(f"An error occurred: {e}")
self.conn.rollback()
src/pgdb/knowledge/pgsqldocstore.py deleted 100644 → 0
import sys
from os import path
# 这里相当于把当前目录添加到pythonpath中
sys.path.append(path.dirname(path.abspath(__file__)))
from typing import List, Union, Dict, Optional
from langchain_community.docstore.base import AddableMixin, Docstore
from k_db import PostgresDB
from .txt_doc_table import TxtDoc
from .vec_txt_table import TxtVector
import json, hashlib, base64
from langchain_core.documents import Document
def str2hash_base64(inp: str) -> str:
# return f"%s" % hash(input)
return base64.b64encode(hashlib.sha1(inp.encode()).digest()).decode()
class PgSqlDocstore(Docstore, AddableMixin):
host: str
dbname: str
username: str
password: str
port: str
'''
说明，重写__getstate__，__setstate__，适用于langchain的序列化存储，基于pickle进行存储。返回数组包含pgsql连接信息。
'''
def __getstate__(self):
return {"host": self.host, "dbname": self.dbname, "username": self.username, "password": self.password,
"port": self.port}
def __setstate__(self, info):
self.__init__(info)
def __init__(self, info: dict, reset: bool = False):
self.host = info["host"]
self.dbname = info["dbname"]
self.username = info["username"]
self.password = info["password"]
self.port = info["port"] if "port" in info else "5432"
self.pgdb = PostgresDB(self.host, self.dbname, self.username, self.password, port=self.port)
self.TXT_DOC = TxtDoc(self.pgdb)
self.VEC_TXT = TxtVector(self.pgdb)
if reset:
self.__sub_init__()
self.TXT_DOC.drop_table()
self.VEC_TXT.drop_table()
self.TXT_DOC.create_table()
self.VEC_TXT.create_table()
def __sub_init__(self):
if not self.pgdb.conn:
self.pgdb.connect()
'''
从本地库中查找向量对应的文本段落，封装成Document返回
'''
def search(self, search: str) -> Union[str, Document]:
if not self.pgdb.conn:
self.__sub_init__()
anwser = self.VEC_TXT.search(search)
content = self.TXT_DOC.search(anwser[0])
meta = json.loads(content[1])
meta.update({"hash": anwser[0]}) # paragraph_id = hash 插入到metadata中，便于后续根据段落查找
if content:
return Document(page_content=content[0], metadata=meta)
else:
return Document()
'''
从本地库中删除向量对应的文本，批量删除
'''
def delete(self, ids: List) -> None:
if not self.pgdb.conn:
self.__sub_init__()
pids = []
for item in ids:
anwser = self.VEC_TXT.search(item)
pids.append(anwser[0])
self.VEC_TXT.delete(ids)
self.TXT_DOC.delete(pids)
'''
向本地库添加向量和文本信息
[vector_id,Document(page_content=问题, metadata=dict(paragraph=段落文本))]
'''
def add(self, texts: Dict[str, Document]) -> None:
# for vec,doc in texts.items():
# paragraph_id = self.TXT_DOC.insert(doc.metadata["paragraph"])
# self.VEC_TXT.insert(vector_id=vec,paragraph_id=paragraph_id,text=doc.page_content)
if not self.pgdb.conn:
self.__sub_init__()
paragraph_hashs = [] # hash,text
paragraph_txts = []
vec_inserts = []
for vec, doc in texts.items():
txt_hash = str2hash_base64(doc.metadata["paragraph"])
print(txt_hash)
vec_inserts.append((vec, doc.page_content, txt_hash))
if txt_hash not in paragraph_hashs:
paragraph_hashs.append(txt_hash)
paragraph = doc.metadata["paragraph"]
doc.metadata.pop("paragraph")
paragraph_txts.append((txt_hash, paragraph, json.dumps(doc.metadata, ensure_ascii=False)))
# print(paragraph_txts)
self.TXT_DOC.insert(paragraph_txts)
self.VEC_TXT.insert(vec_inserts)
class InMemorySecondaryDocstore(Docstore, AddableMixin):
"""Simple in memory docstore in the form of a dict."""
def __init__(self, _dict: Optional[Dict[str, Document]] = None, _sec_dict: Optional[Dict[str, Document]] = None):
"""Initialize with dict."""
self._dict = _dict if _dict is not None else {}
self._sec_dict = _sec_dict if _sec_dict is not None else {}
def add(self, texts: Dict[str, Document]) -> None:
"""Add texts to in memory dictionary.
Args:
texts: dictionary of id -> document.
Returns:
None
"""
overlapping = set(texts).intersection(self._dict)
if overlapping:
raise ValueError(f"Tried to add ids that already exist: {overlapping}")
self._dict = {**self._dict, **texts}
dict1 = {}
dict_sec = {}
for vec, doc in texts.items():
txt_hash = str2hash_base64(doc.metadata["paragraph"])
metadata = doc.metadata
paragraph = metadata.pop('paragraph')
# metadata.update({"paragraph_id":txt_hash})
metadata['paragraph_id'] = txt_hash
dict_sec[txt_hash] = Document(page_content=paragraph, metadata=metadata)
dict1[vec] = Document(page_content=doc.page_content, metadata={'paragraph_id': txt_hash})
self._dict = {**self._dict, **dict1}
self._sec_dict = {**self._sec_dict, **dict_sec}
def delete(self, ids: List) -> None:
"""Deleting IDs from in memory dictionary."""
overlapping = set(ids).intersection(self._dict)
if not overlapping:
raise ValueError(f"Tried to delete ids that does not exist: {ids}")
for _id in ids:
self._sec_dict.pop(self._dict[_id].metadata['paragraph_id'])
self._dict.pop(_id)
def search(self, search: str) -> Union[str, Document]:
"""Search via direct lookup.
Args:
search: id of a document to search for.
Returns:
Document if found, else error message.
"""
if search not in self._dict:
return f"ID {search} not found."
else:
print(self._dict[search].page_content)
return self._sec_dict[self._dict[search].metadata['paragraph_id']]
src/pgdb/knowledge/txt_doc_table.py deleted 100644 → 0
from .k_db import PostgresDB
# paragraph_id BIGSERIAL primary key,
TABLE_TXT_DOC = """
create table txt_doc (
hash varchar(40) primary key,
text text not null,
matadate text
);
"""
TABLE_TXT_DOC_HASH_INDEX = """
CREATE UNIQUE INDEX hash_index ON txt_doc (hash);
"""
# CREATE UNIQUE INDEX idx_name ON your_table (column_name);
class TxtDoc:
def __init__(self, db: PostgresDB) -> None:
self.db = db
def insert(self, texts):
query = f"INSERT INTO txt_doc(hash,text,matadate) VALUES "
args = []
for value in texts:
value = list(value)
query += "(%s,%s,%s),"
args.extend(value)
query = query[:len(query) - 1]
query += f"ON conflict(hash) DO UPDATE SET text = EXCLUDED.text;"
self.db.execute_args(query, args)
def delete(self, ids):
for item in ids:
query = f"delete FROM txt_doc WHERE hash = %s" % item
self.db.execute(query)
def search(self, item):
query = "SELECT text,matadate FROM txt_doc WHERE hash = %s"
self.db.execute_args(query, [item])
answer = self.db.fetchall()
if len(answer) > 0:
return answer[0]
else:
return None
# return Document(page_content=self.db.fetchall()[0][0], metadata=dict(page=self.db.fetchall()[0][1]))
# answer = self.db.fetchall()[0][0]
# return answer
def create_table(self):
query = f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'txt_doc')"
self.db.execute(query)
exists = self.db.fetchall()[0][0]
if not exists:
query = TABLE_TXT_DOC
self.db.execute(query)
# self.db.execute(TABLE_TXT_DOC_HASH_INDEX)
def drop_table(self):
query = f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'txt_doc')"
self.db.execute(query)
exists = self.db.fetchall()[0][0]
if exists:
query = "DROP TABLE txt_doc"
self.db.format(query)
print("drop table txt_doc ok")
def find_like_doc(self,item:str):
query = "select text,matadate FROM txt_doc WHERE matadate like '%"+item+"%' or text like '%"+item+"%' "
self.db.execute(query)
answer = self.db.fetchall()
if len(answer) > 0:
return answer
else:
return None
\ No newline at end of file
src/pgdb/knowledge/vec_txt_table.py deleted 100644 → 0
from .k_db import PostgresDB
TABLE_VEC_TXT = """
CREATE TABLE vec_txt (
vector_id varchar(36) PRIMARY KEY,
text text,
paragraph_id varchar(40) not null
)
"""
# 025a9bee-2eb2-47f5-9722-525e05a0442b
class TxtVector:
def __init__(self, db: PostgresDB) -> None:
self.db = db
def insert(self, vectors):
query = f"INSERT INTO vec_txt(vector_id,text,paragraph_id) VALUES"
args = []
for value in vectors:
value = list(value)
query += "(%s,%s,%s),"
args.extend(value)
query = query[:len(query) - 1]
query += f"ON conflict(vector_id) DO UPDATE SET text = EXCLUDED.text,paragraph_id = EXCLUDED.paragraph_id;"
# query += ";"
self.db.execute_args(query, args)
def delete(self, ids):
for item in ids:
query = f"delete FROM vec_txt WHERE vector_id = '%s'" % (item,)
self.db.execute(query)
def search(self, search: str):
query = f"SELECT paragraph_id,text FROM vec_txt WHERE vector_id = %s"
self.db.execute_args(query, [search])
answer = self.db.fetchall()
# print(answer)
return answer[0]
def create_table(self):
query = f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'vec_txt')"
self.db.execute(query)
exists = self.db.fetchall()[0][0]
if not exists:
query = TABLE_VEC_TXT
self.db.execute(query)
def drop_table(self):
query = f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'vec_txt')"
self.db.execute(query)
exists = self.db.fetchall()[0][0]
if exists:
query = "DROP TABLE vec_txt"
self.db.format(query)
print("drop table vec_txt ok")
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment