mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
XianBW
62a2f7a742
* add CI fix tool to app * perform ruff safe fix * wrap too long lines in prompts.py
616 lines
21 KiB
Python
616 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import multiprocessing as mp
|
|
import re
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Mapping
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import tiktoken
|
|
import yaml
|
|
from azure.ai.formrecognizer import DocumentAnalysisClient
|
|
from azure.core.credentials import AzureKeyCredential
|
|
from jinja2 import Template
|
|
from sklearn.cluster import KMeans
|
|
from sklearn.metrics.pairwise import cosine_similarity
|
|
from sklearn.preprocessing import normalize
|
|
|
|
from core.conf import FincoSettings as Config
|
|
from core.log import FinCoLog
|
|
from oai.llm_utils import APIBackend, create_embedding_with_multiprocessing
|
|
|
|
if TYPE_CHECKING:
|
|
from langchain_core.documents import Document
|
|
|
|
from langchain.document_loaders import PyPDFDirectoryLoader, PyPDFLoader
|
|
|
|
with (Path(__file__).parent / "util_prompt.yaml").open(encoding="utf8") as f:
|
|
UTIL_PROMPT = yaml.safe_load(
|
|
f,
|
|
)
|
|
|
|
|
|
def load_documents_by_langchain(path: Path) -> list:
|
|
"""Load documents from the specified path.
|
|
|
|
Args:
|
|
path (str): The path to the directory or file containing the documents.
|
|
|
|
Returns:
|
|
list: A list of loaded documents.
|
|
"""
|
|
loader = PyPDFDirectoryLoader(str(path), silent_errors=True) if path.is_dir() else PyPDFLoader(str(path))
|
|
return loader.load()
|
|
|
|
|
|
def process_documents_by_langchain(docs: list[Document]) -> dict[str, str]:
|
|
"""Process a list of documents and group them by document name.
|
|
|
|
Args:
|
|
docs (list): A list of documents.
|
|
|
|
Returns:
|
|
dict: A dictionary where the keys are document names and the values are
|
|
the concatenated content of the documents.
|
|
"""
|
|
content_dict = {}
|
|
|
|
for doc in docs:
|
|
doc_name = str(Path(doc.metadata["source"]).resolve())
|
|
doc_content = doc.page_content
|
|
|
|
if doc_name not in content_dict:
|
|
content_dict[str(doc_name)] = doc_content
|
|
else:
|
|
content_dict[str(doc_name)] += doc_content
|
|
|
|
return content_dict
|
|
|
|
|
|
def load_and_process_pdfs_by_langchain(path: Path) -> dict[str, str]:
|
|
return process_documents_by_langchain(load_documents_by_langchain(path))
|
|
|
|
|
|
def load_and_process_one_pdf_by_azure_document_intelligence(
|
|
path: Path,
|
|
key: str,
|
|
endpoint: str,
|
|
) -> str:
|
|
pages = len(PyPDFLoader(str(path)).load())
|
|
document_analysis_client = DocumentAnalysisClient(
|
|
endpoint=endpoint,
|
|
credential=AzureKeyCredential(key),
|
|
)
|
|
|
|
with path.open("rb") as file:
|
|
result = document_analysis_client.begin_analyze_document(
|
|
"prebuilt-document",
|
|
file,
|
|
pages=f"1-{pages}",
|
|
).result()
|
|
return result.content
|
|
|
|
|
|
def load_and_process_pdfs_by_azure_document_intelligence(path: Path) -> dict[str, str]:
|
|
config = Config()
|
|
|
|
assert config.azure_document_intelligence_key is not None
|
|
assert config.azure_document_intelligence_endpoint is not None
|
|
|
|
content_dict = {}
|
|
ab_path = path.resolve()
|
|
if ab_path.is_file():
|
|
assert ".pdf" in ab_path.suffixes, "The file must be a PDF file."
|
|
proc = load_and_process_one_pdf_by_azure_document_intelligence
|
|
content_dict[str(ab_path)] = proc(
|
|
ab_path,
|
|
config.azure_document_intelligence_key,
|
|
config.azure_document_intelligence_endpoint,
|
|
)
|
|
else:
|
|
for file_path in ab_path.rglob("*"):
|
|
if file_path.is_file() and ".pdf" in file_path.suffixes:
|
|
content_dict[str(file_path)] = load_and_process_one_pdf_by_azure_document_intelligence(
|
|
file_path,
|
|
config.azure_document_intelligence_key,
|
|
config.azure_document_intelligence_endpoint,
|
|
)
|
|
return content_dict
|
|
|
|
|
|
def classify_report_from_dict(
|
|
report_dict: Mapping[str, str],
|
|
api: APIBackend,
|
|
input_max_token: int = 128000,
|
|
vote_time: int = 1,
|
|
substrings: tuple[str] = (),
|
|
) -> dict[str, dict[str, str]]:
|
|
"""
|
|
Parameters:
|
|
- report_dict (Dict[str, str]):
|
|
A dictionary where the key is the path of the report (ending with .pdf),
|
|
and the value is either the report content as a string.
|
|
- api (APIBackend): An instance of the APIBackend class.
|
|
- input_max_token (int): Specifying the maximum number of input tokens.
|
|
- vote_time (int): An integer specifying how many times to vote.
|
|
- substrings (list(str)): List of hardcode substrings.
|
|
|
|
Returns:
|
|
- Dict[str, Dict[str, str]]: A dictionary where each key is the path of the report,
|
|
with a single key 'class' and its value being the classification result (0 or 1).
|
|
|
|
"""
|
|
if len(substrings) == 0:
|
|
substrings = (
|
|
"FinCo",
|
|
"金融工程",
|
|
"金工",
|
|
"回测",
|
|
"因子",
|
|
"机器学习",
|
|
"深度学习",
|
|
"量化",
|
|
)
|
|
|
|
res_dict = {}
|
|
classify_prompt = UTIL_PROMPT["classify_system"]
|
|
enc = tiktoken.encoding_for_model("gpt-4-turbo")
|
|
|
|
for key, value in report_dict.items():
|
|
if not key.endswith(".pdf"):
|
|
continue
|
|
file_name = key
|
|
|
|
if isinstance(value, str):
|
|
content = value
|
|
else:
|
|
FinCoLog().warning(f"输入格式不符合要求: {file_name}")
|
|
res_dict[file_name] = {"class": 0}
|
|
continue
|
|
|
|
if not any(substring in content for substring in substrings):
|
|
res_dict[file_name] = {"class": 0}
|
|
else:
|
|
gpt_4_max_token = 128000
|
|
if input_max_token < gpt_4_max_token:
|
|
content = enc.encode(content)
|
|
max_token_1 = max(0, min(len(content), input_max_token) - 1)
|
|
content = enc.decode(content[:max_token_1])
|
|
|
|
vote_list = []
|
|
for _ in range(vote_time):
|
|
user_prompt = content
|
|
system_prompt = classify_prompt
|
|
res = api.build_messages_and_create_chat_completion(
|
|
user_prompt=user_prompt,
|
|
system_prompt=system_prompt,
|
|
json_mode=True,
|
|
)
|
|
try:
|
|
res = json.loads(res)
|
|
vote_list.append(int(res["class"]))
|
|
except json.JSONDecodeError:
|
|
FinCoLog().warning(f"返回值无法解析: {file_name}")
|
|
res_dict[file_name] = {"class": 0}
|
|
count_0 = vote_list.count(0)
|
|
count_1 = vote_list.count(1)
|
|
if max(count_0, count_1) > int(vote_time / 2):
|
|
break
|
|
|
|
result = 1 if count_1 > count_0 else 0
|
|
res_dict[file_name] = {"class": result}
|
|
|
|
return res_dict
|
|
|
|
|
|
def __extract_factors_name_and_desc_from_content(
|
|
content: str,
|
|
) -> dict[str, dict[str, str]]:
|
|
session = APIBackend().build_chat_session(
|
|
session_system_prompt=UTIL_PROMPT["extract_factors_system"],
|
|
)
|
|
|
|
extracted_factor_dict = {}
|
|
current_user_prompt = content
|
|
|
|
for _ in range(10):
|
|
extract_result_resp = session.build_chat_completion(
|
|
user_prompt=current_user_prompt,
|
|
json_mode=False,
|
|
)
|
|
re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
|
|
ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
|
|
try:
|
|
ret_dict = json.loads(ret_json_str)
|
|
parse_success = bool(isinstance(ret_dict, dict)) and "factors" in ret_dict
|
|
except json.JSONDecodeError:
|
|
parse_success = False
|
|
if ret_json_str is None or not parse_success:
|
|
current_user_prompt = (
|
|
"Your response didn't follow the instruction might be wrong json format. Try again."
|
|
)
|
|
else:
|
|
factors = ret_dict["factors"]
|
|
if len(factors) == 0:
|
|
break
|
|
for factor_name, factor_description in factors.items():
|
|
extracted_factor_dict[factor_name] = factor_description
|
|
current_user_prompt = UTIL_PROMPT["extract_factors_follow_user"]
|
|
|
|
return extracted_factor_dict
|
|
|
|
|
|
def __extract_factors_formulation_from_content(
|
|
content: str,
|
|
factor_dict: dict[str, str],
|
|
) -> dict[str, dict[str, str]]:
|
|
factor_dict_df = pd.DataFrame(
|
|
factor_dict.items(),
|
|
columns=["factor_name", "factor_description"],
|
|
)
|
|
|
|
system_prompt = UTIL_PROMPT["extract_factor_formulation_system"]
|
|
current_user_prompt = Template(
|
|
UTIL_PROMPT["extract_factor_formulation_user"],
|
|
).render(report_content=content, factor_dict=factor_dict_df.to_string())
|
|
|
|
session = APIBackend().build_chat_session(session_system_prompt=system_prompt)
|
|
factor_to_formulation = {}
|
|
|
|
for _ in range(10):
|
|
extract_result_resp = session.build_chat_completion(
|
|
user_prompt=current_user_prompt,
|
|
json_mode=False,
|
|
)
|
|
re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
|
|
ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
|
|
try:
|
|
ret_dict = json.loads(ret_json_str)
|
|
parse_success = bool(isinstance(ret_dict, dict))
|
|
except json.JSONDecodeError:
|
|
parse_success = False
|
|
if ret_json_str is None or not parse_success:
|
|
current_user_prompt = (
|
|
"Your response didn't follow the instruction might be wrong json format. Try again."
|
|
)
|
|
else:
|
|
for name, formulation_and_description in ret_dict.items():
|
|
if name in factor_dict:
|
|
factor_to_formulation[name] = formulation_and_description
|
|
if len(factor_to_formulation) != len(factor_dict):
|
|
remain_df = factor_dict_df[~factor_dict_df["factor_name"].isin(factor_to_formulation)]
|
|
current_user_prompt = (
|
|
"Some factors are missing. Please check the following"
|
|
" factors and their descriptions and continue extraction.\n"
|
|
"==========================Remaining factors"
|
|
"==========================\n" + remain_df.to_string()
|
|
)
|
|
else:
|
|
break
|
|
|
|
return factor_to_formulation
|
|
|
|
|
|
def extract_factor_and_formulation_from_one_report(
|
|
content: str,
|
|
) -> dict[str, dict[str, str]]:
|
|
final_factor_dict_to_one_report = {}
|
|
factor_dict = __extract_factors_name_and_desc_from_content(content)
|
|
if len(factor_dict) != 0:
|
|
factor_to_formulation = __extract_factors_formulation_from_content(
|
|
content,
|
|
factor_dict,
|
|
)
|
|
for factor_name in factor_dict:
|
|
final_factor_dict_to_one_report.setdefault(factor_name, {})
|
|
final_factor_dict_to_one_report[factor_name]["description"] = factor_dict[factor_name]
|
|
|
|
# use code to correct _ in formulation
|
|
formulation = factor_to_formulation[factor_name]["formulation"]
|
|
if factor_name in formulation:
|
|
target_factor_name = factor_name.replace("_", r"\_")
|
|
formulation = formulation.replace(factor_name, target_factor_name)
|
|
for variable in factor_to_formulation[factor_name]["variables"]:
|
|
if variable in formulation:
|
|
target_variable = variable.replace("_", r"\_")
|
|
formulation = formulation.replace(variable, target_variable)
|
|
|
|
final_factor_dict_to_one_report[factor_name]["formulation"] = formulation
|
|
final_factor_dict_to_one_report[factor_name]["variables"] = factor_to_formulation[factor_name]["variables"]
|
|
|
|
return final_factor_dict_to_one_report
|
|
|
|
|
|
def extract_factors_from_report_dict_and_classify_result(
|
|
report_dict: dict[str, str],
|
|
useful_no_dict: dict[str, dict[str, str]],
|
|
n_proc: int = 11,
|
|
) -> dict[str, dict[str, dict[str, str]]]:
|
|
useful_report_dict = {}
|
|
for key, value in useful_no_dict.items():
|
|
if isinstance(value, dict):
|
|
if int(value.get("class")) == 1:
|
|
useful_report_dict[key] = report_dict[key]
|
|
else:
|
|
FinCoLog().warning(f"输入格式不符合要求: {key}")
|
|
|
|
final_report_factor_dict = {}
|
|
# for file_name, content in useful_report_dict.items():
|
|
# final_report_factor_dict.setdefault(file_name, {})
|
|
# final_report_factor_dict[
|
|
# file_name
|
|
# ] = extract_factor_and_formulation_from_one_report(content)
|
|
|
|
while len(final_report_factor_dict) != len(useful_report_dict):
|
|
pool = mp.Pool(n_proc)
|
|
pool_result_list = []
|
|
file_names = []
|
|
for file_name, content in useful_report_dict.items():
|
|
if file_name in final_report_factor_dict:
|
|
continue
|
|
file_names.append(file_name)
|
|
pool_result_list.append(
|
|
pool.apply_async(
|
|
extract_factor_and_formulation_from_one_report,
|
|
(content,),
|
|
),
|
|
)
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
for index, result in enumerate(pool_result_list):
|
|
if result.get is not None:
|
|
file_name = file_names[index]
|
|
final_report_factor_dict.setdefault(file_name, {})
|
|
final_report_factor_dict[file_name] = result.get()
|
|
FinCoLog().info(f"已经完成{len(final_report_factor_dict)}个报告的因子提取")
|
|
|
|
return final_report_factor_dict
|
|
|
|
|
|
def check_factor_dict_viability_simulate_json_mode(
|
|
factor_df_string: str,
|
|
) -> dict[str, dict[str, str]]:
|
|
session = APIBackend().build_chat_session(
|
|
session_system_prompt=UTIL_PROMPT["factor_viability_system"],
|
|
)
|
|
current_user_prompt = factor_df_string
|
|
|
|
for _ in range(10):
|
|
extract_result_resp = session.build_chat_completion(
|
|
user_prompt=current_user_prompt,
|
|
json_mode=False,
|
|
)
|
|
re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
|
|
ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
|
|
try:
|
|
ret_dict = json.loads(ret_json_str)
|
|
parse_success = bool(isinstance(ret_dict, dict))
|
|
except json.JSONDecodeError:
|
|
parse_success = False
|
|
if ret_json_str is None or not parse_success:
|
|
current_user_prompt = (
|
|
"Your response didn't follow the instruction might be wrong json format. Try again."
|
|
)
|
|
else:
|
|
return ret_dict
|
|
return {}
|
|
|
|
|
|
def check_factor_dict_viability(
|
|
factor_dict: dict[str, dict[str, str]],
|
|
) -> dict[str, dict[str, str]]:
|
|
factor_viability_dict = {}
|
|
|
|
factor_df = pd.DataFrame(factor_dict).T
|
|
factor_df.index.names = ["factor_name"]
|
|
|
|
while factor_df.shape[0] > 0:
|
|
pool = mp.Pool(8)
|
|
|
|
result_list = []
|
|
for i in range(0, factor_df.shape[0], 50):
|
|
target_factor_df_string = factor_df.iloc[i : i + 50, :].to_string()
|
|
|
|
result_list.append(
|
|
pool.apply_async(
|
|
check_factor_dict_viability_simulate_json_mode,
|
|
(target_factor_df_string,),
|
|
),
|
|
)
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
for result in result_list:
|
|
respond = result.get()
|
|
for factor_name, viability in respond.items():
|
|
factor_viability_dict[factor_name] = viability
|
|
|
|
factor_df = factor_df[~factor_df.index.isin(factor_viability_dict)]
|
|
|
|
return factor_viability_dict
|
|
|
|
|
|
def check_factor_duplication_simulate_json_mode(
|
|
factor_df: pd.DataFrame,
|
|
) -> list[list[str]]:
|
|
session = APIBackend().build_chat_session(
|
|
session_system_prompt=UTIL_PROMPT["factor_duplicate_system"],
|
|
)
|
|
current_user_prompt = factor_df.to_string()
|
|
|
|
generated_duplicated_groups = []
|
|
for _ in range(20):
|
|
extract_result_resp = session.build_chat_completion(
|
|
user_prompt=current_user_prompt,
|
|
json_mode=False,
|
|
)
|
|
re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
|
|
ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
|
|
try:
|
|
ret_dict = json.loads(ret_json_str)
|
|
parse_success = bool(isinstance(ret_dict, list))
|
|
except json.JSONDecodeError:
|
|
parse_success = False
|
|
if ret_json_str is None or not parse_success:
|
|
current_user_prompt = (
|
|
"Your previous response didn't follow"
|
|
" the instruction might be wrong json"
|
|
" format. Try reducing the factors."
|
|
)
|
|
elif len(ret_dict) == 0:
|
|
return generated_duplicated_groups
|
|
else:
|
|
generated_duplicated_groups.extend(ret_dict)
|
|
current_user_prompt = (
|
|
"Continue to extract duplicated"
|
|
" groups. If no more duplicated group"
|
|
" found please respond empty dict."
|
|
)
|
|
return generated_duplicated_groups
|
|
|
|
|
|
def kmeans_embeddings(embeddings: np.ndarray, k: int = 20) -> list[list[str]]:
|
|
x_normalized = normalize(embeddings)
|
|
|
|
kmeans = KMeans(
|
|
n_clusters=k,
|
|
init="random",
|
|
max_iter=100,
|
|
n_init=10,
|
|
random_state=42,
|
|
)
|
|
|
|
# KMeans算法使用欧氏距离, 我们需要自定义一个函数来找到最相似的簇中心
|
|
def find_closest_cluster_cosine_similarity(
|
|
data: np.ndarray,
|
|
centroids: np.ndarray,
|
|
) -> np.ndarray:
|
|
similarity = cosine_similarity(data, centroids)
|
|
return np.argmax(similarity, axis=1)
|
|
|
|
# 初始化簇中心
|
|
rng = np.random.default_rng()
|
|
centroids = rng.choice(x_normalized, size=k, replace=False)
|
|
|
|
# 迭代直到收敛或达到最大迭代次数
|
|
for _ in range(kmeans.max_iter):
|
|
# 分配样本到最近的簇中心
|
|
closest_clusters = find_closest_cluster_cosine_similarity(
|
|
x_normalized,
|
|
centroids,
|
|
)
|
|
|
|
# 更新簇中心
|
|
new_centroids = np.array(
|
|
[x_normalized[closest_clusters == i].mean(axis=0) for i in range(k)],
|
|
)
|
|
new_centroids = normalize(new_centroids) # 归一化新的簇中心
|
|
|
|
# 检查簇中心是否发生变化
|
|
if np.allclose(centroids, new_centroids):
|
|
break
|
|
|
|
centroids = new_centroids
|
|
|
|
clusters = find_closest_cluster_cosine_similarity(x_normalized, centroids)
|
|
cluster_to_index = {}
|
|
for index, cluster in enumerate(clusters):
|
|
cluster_to_index.setdefault(cluster, []).append(index)
|
|
return sorted(
|
|
cluster_to_index.values(),
|
|
key=lambda x: len(x),
|
|
reverse=True,
|
|
)
|
|
|
|
|
|
def deduplicate_factor_dict(factor_dict: dict[str, dict[str, str]]) -> list[list[str]]:
|
|
factor_df = pd.DataFrame(factor_dict).T
|
|
factor_df.index.names = ["factor_name"]
|
|
|
|
factor_names = sorted(factor_dict)
|
|
|
|
factor_name_to_full_str = {}
|
|
for factor_name in factor_dict:
|
|
description = factor_dict[factor_name]["description"]
|
|
formulation = factor_dict[factor_name]["formulation"]
|
|
variables = factor_dict[factor_name]["variables"]
|
|
factor_name_to_full_str[
|
|
factor_name
|
|
] = f"""Factor name: {factor_name}
|
|
Factor description: {description}
|
|
Factor formulation: {formulation}
|
|
Factor variables: {variables}
|
|
"""
|
|
|
|
full_str_list = [factor_name_to_full_str[factor_name] for factor_name in factor_names]
|
|
embeddings = create_embedding_with_multiprocessing(full_str_list)
|
|
|
|
target_k = None
|
|
if len(full_str_list) < Config().max_input_duplicate_factor_group:
|
|
kmeans_index_group = [list(range(len(full_str_list)))]
|
|
target_k = 1
|
|
else:
|
|
for k in range(
|
|
len(full_str_list) // Config().max_input_duplicate_factor_group,
|
|
30,
|
|
):
|
|
kmeans_index_group = kmeans_embeddings(embeddings=embeddings, k=k)
|
|
if len(kmeans_index_group[0]) < Config().max_input_duplicate_factor_group:
|
|
target_k = k
|
|
FinCoLog().info(f"K-means group number: {k}")
|
|
break
|
|
factor_name_groups = [[factor_names[index] for index in index_group] for index_group in kmeans_index_group]
|
|
|
|
duplication_names_list = []
|
|
|
|
pool = mp.Pool(target_k)
|
|
result_list = []
|
|
result_list = [
|
|
pool.apply_async(
|
|
check_factor_duplication_simulate_json_mode,
|
|
(factor_df.loc[factor_name_group, :],),
|
|
)
|
|
for factor_name_group in factor_name_groups
|
|
]
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
for result in result_list:
|
|
deduplication_factor_names_list = result.get()
|
|
for deduplication_factor_names in deduplication_factor_names_list:
|
|
filter_factor_names = [
|
|
factor_name for factor_name in set(deduplication_factor_names) if factor_name in factor_dict
|
|
]
|
|
if len(filter_factor_names) > 1:
|
|
duplication_names_list.append(filter_factor_names)
|
|
|
|
return duplication_names_list
|
|
|
|
|
|
def deduplicate_factors_several_times(
|
|
factor_dict: dict[str, dict[str, str]],
|
|
) -> list[list[str]]:
|
|
final_duplication_names_list = []
|
|
current_round_factor_dict = factor_dict
|
|
for _ in range(10):
|
|
duplication_names_list = deduplicate_factor_dict(current_round_factor_dict)
|
|
|
|
new_round_names = []
|
|
for duplication_names in duplication_names_list:
|
|
if len(duplication_names) < Config().max_output_duplicate_factor_group:
|
|
final_duplication_names_list.append(duplication_names)
|
|
else:
|
|
new_round_names.extend(duplication_names)
|
|
|
|
if len(new_round_names) != 0:
|
|
current_round_factor_dict = {factor_name: factor_dict[factor_name] for factor_name in new_round_names}
|
|
else:
|
|
return final_duplication_names_list
|
|
return []
|