|
|
# agent/schema/schema_utils.py
|
|
|
import jieba
|
|
|
import numpy as np
|
|
|
from collections import Counter
|
|
|
from config import *
|
|
|
from agent.utils import clean_text, find_dict_by_element
|
|
|
|
|
|
def parse_table_structures(input_text):
|
|
|
"""
|
|
|
用于解析输入文本并提取表结构信息，返回一个字典，其中表名作为键，表结构作为值
|
|
|
Parse the input text to extract table structures.
|
|
|
|
|
|
The format is expected as pairs: "table_name === table_structure".
|
|
|
|
|
|
Parameters:
|
|
|
input_text (str): The raw text containing table structures.
|
|
|
|
|
|
Returns:
|
|
|
tables_dict (dict): A dictionary where keys are table names and
|
|
|
values are the associated table structures.
|
|
|
"""
|
|
|
tables_text = input_text.split('===')[1:]
|
|
|
tables_dict = {tables_text[i]: tables_text[i + 1] for i in range(0, len(tables_text), 2)}
|
|
|
return tables_dict
|
|
|
|
|
|
def map_chinese_to_english_tables(chinese_names, english_names):
|
|
|
"""
|
|
|
将中文的表名映射到对应的英文表名
|
|
|
Map Chinese table names to their corresponding English table names.
|
|
|
For each Chinese name, there is a matching English name
|
|
|
(case-insensitive comparison).
|
|
|
|
|
|
Parameters:
|
|
|
chinese_names (list): A list of Chinese table names.
|
|
|
english_names (list): A list of English table names.
|
|
|
|
|
|
Returns:
|
|
|
name_map (dict): A dictionary mapping Chinese table names to English table names.
|
|
|
"""
|
|
|
name_map = {}
|
|
|
for cname in chinese_names:
|
|
|
# Find the corresponding English name (case-insensitive match)
|
|
|
english_match = [en for en in english_names if str(en).lower() == cname.lower()][0]
|
|
|
name_map[cname] = english_match
|
|
|
return name_map
|
|
|
|
|
|
def find_value_in_list_of_dicts(dict_list, key_to_match, value_to_match, key_to_return):
|
|
|
"""
|
|
|
在字典列表中查找满足条件的字典，并返回其中指定键的值。
|
|
|
Search through a list of dictionaries and find the first dictionary where
|
|
|
the value of key_to_match equals value_to_match, then return the value
|
|
|
associated with key_to_return.
|
|
|
|
|
|
Parameters:
|
|
|
dict_list (list): A list of dictionaries to search through.
|
|
|
key_to_match (str): The key whose value we want to match.
|
|
|
value_to_match (str): The value we are looking for.
|
|
|
key_to_return (str): The key whose value we want to return.
|
|
|
|
|
|
Returns:
|
|
|
(str): The value associated with key_to_return in the matching dictionary,
|
|
|
or an empty string if no match is found.
|
|
|
"""
|
|
|
for dictionary in dict_list:
|
|
|
if dictionary.get(key_to_match) == value_to_match:
|
|
|
return dictionary.get(key_to_return)
|
|
|
return ''
|
|
|
|
|
|
def filter_table_comments(question, table_comments):
|
|
|
"""
|
|
|
根据输入问题从表注释列表中筛选出与问题相关的注释。
|
|
|
Filter a list of table comments based on the given question.
|
|
|
Uses jieba for segmentation and removes stopwords, returning only comments
|
|
|
that contain at least one of the segmented keywords.
|
|
|
|
|
|
Parameters:
|
|
|
question (str): The question text.
|
|
|
table_comments (list): A list of comment strings to filter.
|
|
|
|
|
|
Returns:
|
|
|
filtered_comments (list): Filtered list of comments.
|
|
|
"""
|
|
|
stopwords = ['？', '有', '的', '多少', '人', '（', '）']
|
|
|
seg_list = list(jieba.cut(question, cut_all=False))
|
|
|
filtered_seg_list = [word for word in seg_list if word not in stopwords]
|
|
|
|
|
|
filtered_comments = []
|
|
|
for comment in table_comments:
|
|
|
if any(keyword in comment for keyword in filtered_seg_list):
|
|
|
filtered_comments.append(comment)
|
|
|
return filtered_comments
|
|
|
|
|
|
def get_table_schema(question=''):
|
|
|
"""
|
|
|
获取表格的结构信息以及字段的注释
|
|
|
Retrieve table schemas along with optional filtered field comments.
|
|
|
If a question is provided, the comments will be filtered based on
|
|
|
question keywords.
|
|
|
|
|
|
The function:
|
|
|
1. Maps Chinese table names to English table names.
|
|
|
2. For each table, retrieves its structure and finds associated comments.
|
|
|
3. If a question is provided, filter the comments based on keywords extracted from the question.
|
|
|
|
|
|
Parameters:
|
|
|
question (str): The question text. If empty, no filtering is performed.
|
|
|
|
|
|
Returns:
|
|
|
table_maps (list): A list of dictionaries, each containing table schema information.
|
|
|
{
|
|
|
'数据表名': EnglishTableName,
|
|
|
'数据表结构': TableStructure,
|
|
|
'字段注释': FilteredComments (optional if question is provided)
|
|
|
}
|
|
|
"""
|
|
|
if QUESTION_TYPE == "全股":
|
|
|
parsed_tables = parse_table_structures(input_text_all)
|
|
|
database_L = database_L_all
|
|
|
database_table_en = database_table_en_all
|
|
|
elif QUESTION_TYPE == "港股":
|
|
|
parsed_tables = parse_table_structures(input_text_hk)
|
|
|
database_L = database_L_hk
|
|
|
database_table_en = database_table_en_hk
|
|
|
elif QUESTION_TYPE == "美股":
|
|
|
parsed_tables = parse_table_structures(input_text_us)
|
|
|
database_L = database_L_us
|
|
|
database_table_en = database_table_en_us
|
|
|
elif QUESTION_TYPE == "A股":
|
|
|
parsed_tables = parse_table_structures(input_text_cn)
|
|
|
database_L = database_L_cn
|
|
|
database_table_en = database_table_en_cn
|
|
|
else:
|
|
|
parsed_tables = parse_table_structures(input_text_all)
|
|
|
database_L = database_L_all
|
|
|
database_table_en = database_table_en_all
|
|
|
|
|
|
# Clean up keys and values
|
|
|
cleaned_tables = {
|
|
|
k.replace(' ', '').replace('表结构', ''): v.replace('--', '')
|
|
|
for k, v in parsed_tables.items()
|
|
|
}
|
|
|
|
|
|
# List of Chinese table names (keys)
|
|
|
chinese_table_names = list(cleaned_tables.keys())
|
|
|
|
|
|
name_map = map_chinese_to_english_tables(chinese_table_names, database_table_en)
|
|
|
|
|
|
table_maps = []
|
|
|
for cname, structure in cleaned_tables.items():
|
|
|
english_name = name_map.get(cname)
|
|
|
comments = find_value_in_list_of_dicts(database_L, '数据表名', english_name, '注释')
|
|
|
|
|
|
if question == '':
|
|
|
# No filtering, just return table name and structure
|
|
|
table_map = {
|
|
|
'数据表名': english_name,
|
|
|
'数据表结构': structure
|
|
|
}
|
|
|
else:
|
|
|
# Filter comments based on question
|
|
|
filtered_comments = filter_table_comments(question, comments)
|
|
|
table_map = {
|
|
|
'数据表名': english_name,
|
|
|
'数据表结构': structure,
|
|
|
'字段注释': filtered_comments
|
|
|
}
|
|
|
|
|
|
table_maps.append(table_map)
|
|
|
return table_maps
|
|
|
|
|
|
def to_get_question_columns(question):
|
|
|
"""
|
|
|
Given a question (string) and a global variable database_L_zh (list of dicts),
|
|
|
find 列名 that correspond to 列名中文描述 mentioned in the question.
|
|
|
|
|
|
If any matching columns are found, return a message instructing the user to
|
|
|
use these column names directly for data querying. If none are found, return an empty string.
|
|
|
|
|
|
Parameters:
|
|
|
question (str): The input question text.
|
|
|
|
|
|
Returns:
|
|
|
str: A message with identified column names or an empty string if none found.
|
|
|
"""
|
|
|
database_L_zh = database_L_zh_all
|
|
|
if len(QUESTION_TYPE_LIST) == 0:
|
|
|
database_L_zh = database_L_zh_all
|
|
|
elif len(QUESTION_TYPE_LIST) >= 2:
|
|
|
database_L_zh = database_L_zh_all
|
|
|
elif "港股" in QUESTION_TYPE_LIST:
|
|
|
database_L_zh = database_L_zh_hk
|
|
|
elif "美股" in QUESTION_TYPE_LIST:
|
|
|
database_L_zh = database_L_zh_us
|
|
|
elif "A股" in QUESTION_TYPE_LIST:
|
|
|
database_L_zh = database_L_zh_cn
|
|
|
else:
|
|
|
database_L_zh = database_L_zh_all
|
|
|
|
|
|
L_num = []
|
|
|
for items in database_L_zh:
|
|
|
L_num += items['列名中文描述']
|
|
|
|
|
|
# Get unique column descriptions
|
|
|
L_num_new = [item for item, count in Counter(L_num).items() if count == 1]
|
|
|
|
|
|
# Drop NaN if any
|
|
|
series_num = pd.Series(L_num_new)
|
|
|
L_num_new = list(series_num.dropna())
|
|
|
|
|
|
# Remove known irrelevant items
|
|
|
irrelevant_items = ['年度', '占比']
|
|
|
for irr in irrelevant_items:
|
|
|
if irr in L_num_new:
|
|
|
L_num_new.remove(irr)
|
|
|
|
|
|
matched_columns = []
|
|
|
for col_descs in L_num_new:
|
|
|
col_desc_another = items_another[col_descs]
|
|
|
for col_desc in col_desc_another:
|
|
|
# Check if the column description or its cleaned version appears in the question
|
|
|
if col_desc in question or clean_text(col_desc) in question:
|
|
|
L_dict = find_dict_by_element(database_L_zh, col_descs)
|
|
|
if not L_dict:
|
|
|
break
|
|
|
# Create a mapping from Chinese description to English column name
|
|
|
dict_zip = dict(zip(L_dict[0]['列名中文描述'], L_dict[0]['列名']))
|
|
|
column_name = dict_zip[col_descs]
|
|
|
data_table = L_dict[0]['数据表名']
|
|
|
|
|
|
matched_columns.append({
|
|
|
'数据库表': data_table,
|
|
|
'列名': column_name,
|
|
|
'列名中文含义': col_descs
|
|
|
})
|
|
|
break
|
|
|
|
|
|
if matched_columns:
|
|
|
return f"已获得一部分数据库列名{matched_columns}，请充分利用获得的列名直接查询数据。"
|
|
|
else:
|
|
|
return ''
|
|
|
|
|
|
def get_table_schema_with_emb(tables_name=[], question='', threshold=0.35, print_all=True):
|
|
|
# tables_name是英文名的list
|
|
|
if len(QUESTION_TYPE_LIST) == 0:
|
|
|
parsed_tables = parse_table_structures(input_text_all)
|
|
|
database_L = database_L_all
|
|
|
database_table_en = database_table_en_all
|
|
|
elif len(QUESTION_TYPE_LIST) >= 2:
|
|
|
parsed_tables = parse_table_structures(input_text_all)
|
|
|
database_L = database_L_all
|
|
|
database_table_en = database_table_en_all
|
|
|
elif "港股" in QUESTION_TYPE_LIST:
|
|
|
parsed_tables = parse_table_structures(input_text_hk)
|
|
|
database_L = database_L_hk
|
|
|
database_table_en = database_table_en_hk
|
|
|
elif "美股" in QUESTION_TYPE_LIST:
|
|
|
parsed_tables = parse_table_structures(input_text_us)
|
|
|
database_L = database_L_us
|
|
|
database_table_en = database_table_en_us
|
|
|
elif "A股" in QUESTION_TYPE_LIST:
|
|
|
parsed_tables = parse_table_structures(input_text_cn)
|
|
|
database_L = database_L_cn
|
|
|
database_table_en = database_table_en_cn
|
|
|
else:
|
|
|
parsed_tables = parse_table_structures(input_text_all)
|
|
|
database_L = database_L_all
|
|
|
database_table_en = database_table_en_all
|
|
|
|
|
|
# Clean up keys and values
|
|
|
cleaned_tables = {
|
|
|
k.replace(' ', '').replace('表结构', ''): v.replace('--', '')
|
|
|
for k, v in parsed_tables.items()
|
|
|
}
|
|
|
|
|
|
columns_to_clean = ["column_description"]
|
|
|
time_columns = ["TradingDay", "EndDate", "InfoPublDate"]
|
|
|
remove_columns = ['InitialInfoPublDate', 'XGRQ','InsertTime','UpdateTime']
|
|
|
table_maps = []
|
|
|
highest_score_list = []
|
|
|
if tables_name == [] or question == '':
|
|
|
return None
|
|
|
for table_name in tables_name:
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"---------->table_name:{table_name}")
|
|
|
for english_name, structure in cleaned_tables.items():
|
|
|
if english_name.lower() == table_name.lower():
|
|
|
# print(f"------------threshold:{threshold}")
|
|
|
filtered_comments = []
|
|
|
df_tmp = df_all2[df_all2["table_name"]==table_name.split(".")[1]]
|
|
|
imp_matched_columns, matched_columns_with_scores = extract_matching_columns_with_similarity(
|
|
|
df_tmp, question, threshold, MAX_TOP_COLUMNS
|
|
|
)
|
|
|
if matched_columns_with_scores == [] and imp_matched_columns == []:
|
|
|
# print(f"------------pass")
|
|
|
pass
|
|
|
else:
|
|
|
# print(f"------------no pass")
|
|
|
# 过滤掉不需要的列
|
|
|
remove_columns = ["InitialInfoPublDate", "XGRQ", "InsertTime", "UpdateTime"]
|
|
|
matched_columns = [
|
|
|
item["column"] for item in matched_columns_with_scores if item["column"] not in remove_columns
|
|
|
]
|
|
|
matched_columns = [column for column in matched_columns if column not in remove_columns]
|
|
|
highest_score_item = max(matched_columns_with_scores, key=lambda x: x["similarity_score"])
|
|
|
highest_score = highest_score_item["similarity_score"]
|
|
|
if print_all:
|
|
|
pass
|
|
|
# print(f"------------table_name:{table_name}")
|
|
|
# print(f"------------imp_matched_columns:{imp_matched_columns}, matched_columns_with_scores:{matched_columns_with_scores}")
|
|
|
else:
|
|
|
if (highest_score > threshold) or imp_matched_columns != []:
|
|
|
pass
|
|
|
# print(f"------------table_name:{table_name}")
|
|
|
# print(f"------------imp_matched_columns:{imp_matched_columns}, matched_columns_with_scores:{matched_columns_with_scores}")
|
|
|
filtered_structure = filter_structure(structure, matched_columns, time_columns)
|
|
|
filtered_comments = extract_column_descriptions(filtered_structure, df_tmp)
|
|
|
table_map = {
|
|
|
'数据表名': table_name,
|
|
|
'数据表结构': filtered_structure,
|
|
|
'部分字段注释补充': filtered_comments
|
|
|
}
|
|
|
table_maps.append(table_map)
|
|
|
highest_score_list.append(highest_score)
|
|
|
if not table_maps or not highest_score_list:
|
|
|
return [], []
|
|
|
|
|
|
# 将 table_maps 和 highest_score_list 按分数排序
|
|
|
combined_list = list(zip(table_maps, highest_score_list))
|
|
|
sorted_combined_list = sorted(combined_list, key=lambda x: x[1], reverse=True)
|
|
|
table_maps, highest_score_list = zip(*sorted_combined_list)
|
|
|
table_maps = list(table_maps)
|
|
|
highest_score_list = list(highest_score_list)
|
|
|
return table_maps, highest_score_list
|
|
|
|
|
|
def deep_find_tables(all_tables_name_list=[], question='', threshold=0.6, print_all=False, top_table_n=5):
|
|
|
# tables_name是英文名的list
|
|
|
if len(QUESTION_TYPE_LIST) == 0:
|
|
|
df_tmp = df_all
|
|
|
elif len(QUESTION_TYPE_LIST) >= 2:
|
|
|
df_tmp = df_all
|
|
|
elif "港股" in QUESTION_TYPE_LIST:
|
|
|
df_tmp = df_hk
|
|
|
elif "美股" in QUESTION_TYPE_LIST:
|
|
|
df_tmp = df_us
|
|
|
elif "A股" in QUESTION_TYPE_LIST:
|
|
|
df_tmp = df_cn
|
|
|
else:
|
|
|
df_tmp = df_all
|
|
|
if all_tables_name_list==[] or question=='':
|
|
|
return []
|
|
|
all_table_list = df_tmp['库表名英文'].values.tolist()
|
|
|
all_table_list = [table for table in all_table_list if table not in all_tables_name_list]
|
|
|
all_table_list = list(set(all_table_list))
|
|
|
table_maps, highest_score_list = get_table_schema_with_emb(all_table_list, question, threshold, print_all)
|
|
|
return table_maps[:top_table_n]
|
|
|
|
|
|
# %% [code] {"execution":{"iopub.status.busy":"2025-03-07T09:14:36.671853Z","iopub.execute_input":"2025-03-07T09:14:36.672272Z","iopub.status.idle":"2025-03-07T09:14:36.693550Z","shell.execute_reply.started":"2025-03-07T09:14:36.672191Z","shell.execute_reply":"2025-03-07T09:14:36.692326Z"},"jupyter":{"source_hidden":true}}
|
|
|
def find_table_with_emb(all_tables_name_list=[], question='', use_local_vectors=USE_LOCAL_VECTORS):
|
|
|
"""
|
|
|
判断后返回最高分的表（始终返回1个）：
|
|
|
1. 则对 df_tmp 中所有唯一的库表进行批量比较，并记录每个表的相似度。
|
|
|
2. 最后合并候选结果，从中选取相似度最高的表返回（返回列表中只含1个表）。
|
|
|
"""
|
|
|
# 根据 QUESTION_TYPE_LIST 选择对应的数据集
|
|
|
if len(QUESTION_TYPE_LIST) == 0:
|
|
|
df_tmp = df_all
|
|
|
embeddings_path = df_all_embeddings_path
|
|
|
elif len(QUESTION_TYPE_LIST) >= 2:
|
|
|
df_tmp = df_all
|
|
|
embeddings_path = df_all_embeddings_path
|
|
|
elif "港股" in QUESTION_TYPE_LIST:
|
|
|
df_tmp = df_hk
|
|
|
embeddings_path = df_hk_embeddings_path
|
|
|
elif "美股" in QUESTION_TYPE_LIST:
|
|
|
df_tmp = df_us
|
|
|
embeddings_path = df_us_embeddings_path
|
|
|
elif "A股" in QUESTION_TYPE_LIST:
|
|
|
df_tmp = df_cn
|
|
|
embeddings_path = df_cn_embeddings_path
|
|
|
else:
|
|
|
df_tmp = df_all
|
|
|
embeddings_path = df_all_embeddings_path
|
|
|
|
|
|
if not all_tables_name_list or question == '':
|
|
|
return []
|
|
|
|
|
|
candidate_scores = []
|
|
|
df_unique = df_tmp[['库表名英文', 'table_describe']].drop_duplicates()
|
|
|
candidate_ids_deep = df_unique['库表名英文'].tolist()
|
|
|
candidate_texts_deep = df_unique['table_describe'].tolist()
|
|
|
|
|
|
if use_local_vectors:
|
|
|
texts_for_batch_deep = [question]
|
|
|
similarity_scores_deep = calculate_similarity(texts_for_batch_deep, local_vectors=embeddings_path)
|
|
|
else:
|
|
|
texts_for_batch_deep = [question] + candidate_texts_deep
|
|
|
similarity_scores_deep = calculate_similarity(texts_for_batch_deep)
|
|
|
|
|
|
for table, sim in zip(candidate_ids_deep, similarity_scores_deep):
|
|
|
candidate_scores.append((table, sim))
|
|
|
# 如果没有候选表，则返回空列表
|
|
|
if not candidate_scores:
|
|
|
return []
|
|
|
|
|
|
# 从所有候选中选取相似度最高的表（不论是否达到 threshold）
|
|
|
best_table, best_sim = max(candidate_scores, key=lambda x: x[1])
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"Best table: {best_table} with similarity: {best_sim}")
|
|
|
return [best_table]
|
|
|
|
|
|
def extract_matching_columns_with_similarity(df, question, threshold, top_n=25):
|
|
|
"""
|
|
|
根据question提取匹配列：
|
|
|
1. 对于每行的字段（如 column_description 与 注释），先做全量匹配，
|
|
|
如果文本完全包含在question中，则直接赋分1.0。
|
|
|
2. 同时收集所有候选文本，最后批量调用 calculate_similarity，
|
|
|
将 question 与所有候选文本比较，得到相似度分数（只保留达到阈值的）。
|
|
|
3. 同一列可能对应多个候选文本，最终保留相似度最高的得分。
|
|
|
返回：
|
|
|
- imp_matched_columns：全量匹配（得分1.0）的列名列表
|
|
|
- topn_matched_columns：满足阈值的列及其相似度分数（按分数降序排列，最多 top_n 个）
|
|
|
"""
|
|
|
imp_matched_columns_with_scores = {} # 全量匹配得分为1.0的列
|
|
|
matched_columns_with_scores = {} # 语义匹配得分（可能多次出现，取最高分）
|
|
|
|
|
|
candidate_texts = [] # 收集待计算相似度的候选文本
|
|
|
candidate_ids = [] # 对应候选文本所属的列名
|
|
|
|
|
|
# 遍历DataFrame中的每一行
|
|
|
for _, row in df.iterrows():
|
|
|
col_name = row["column_name"]
|
|
|
|
|
|
# 处理 column_description 字段
|
|
|
if isinstance(row["column_description"], str):
|
|
|
col_desc = row["column_description"]
|
|
|
# 全量匹配检查：如果候选文本出现在question中，则直接记1.0分
|
|
|
if col_desc in question:
|
|
|
imp_matched_columns_with_scores[col_name] = 1.0
|
|
|
matched_columns_with_scores[col_name] = 1.0
|
|
|
# 无论全量匹配与否，都加入候选列表进行语义匹配
|
|
|
else:
|
|
|
candidate_texts.append(col_desc)
|
|
|
candidate_ids.append(col_name)
|
|
|
|
|
|
# 处理 注释 字段（假设多个注释以“，”分隔）
|
|
|
if isinstance(row["注释"], str):
|
|
|
words = row["注释"].split("，")
|
|
|
for word in words:
|
|
|
if word in question:
|
|
|
imp_matched_columns_with_scores[col_name] = 1.0
|
|
|
matched_columns_with_scores[col_name] = 1.0
|
|
|
# candidate_texts.append(word)
|
|
|
# candidate_ids.append(col_name)
|
|
|
|
|
|
# 批量计算语义相似度：将 question 与所有候选文本一次性比较
|
|
|
if candidate_texts:
|
|
|
# 构造输入列表：第一个文本为 question，其余为所有候选文本
|
|
|
texts_for_batch = [question] + candidate_texts
|
|
|
similarity_scores = calculate_similarity(texts_for_batch)
|
|
|
# 遍历每个候选文本的相似度
|
|
|
for idx, sim in enumerate(similarity_scores):
|
|
|
if sim >= threshold:
|
|
|
col = candidate_ids[idx]
|
|
|
# 同一列可能有多个候选文本，保留相似度最高的得分
|
|
|
matched_columns_with_scores[col] = max(matched_columns_with_scores.get(col, 0), sim)
|
|
|
|
|
|
# 合并全量匹配和语义匹配的结果
|
|
|
all_matched_columns = {**matched_columns_with_scores, **imp_matched_columns_with_scores}
|
|
|
# 按相似度分数降序排列（注意：由于全量匹配的分数为1.0，通常会排在最前面）
|
|
|
unique_matched_columns = {col: score for col, score in sorted(all_matched_columns.items(), key=lambda x: -x[1])}
|
|
|
|
|
|
# 提取全量匹配列
|
|
|
imp_matched_columns = [col for col, score in unique_matched_columns.items() if score == 1.0]
|
|
|
|
|
|
# 筛选出满足阈值的列，并构造输出格式（字典列表），按分数降序排序
|
|
|
matched_columns = sorted(
|
|
|
[{"column": col, "similarity_score": score} for col, score in unique_matched_columns.items() if score >= threshold],
|
|
|
key=lambda x: x["similarity_score"],
|
|
|
reverse=True
|
|
|
)
|
|
|
topn_matched_columns = matched_columns[:top_n]
|
|
|
return imp_matched_columns, topn_matched_columns
|
|
|
|
|
|
def filter_structure(structure, matched_columns, time_columns):
|
|
|
# 分割表头和数据部分
|
|
|
sections = structure.split("\n\n", 1)
|
|
|
header = sections[0] # 表头部分
|
|
|
data_rows = sections[1] if len(sections) > 1 else "" # 数据部分
|
|
|
|
|
|
# 条件检查函数
|
|
|
def satisfies_conditions(row):
|
|
|
row_fields = row.split() # 假设字段之间是用空格分隔的
|
|
|
# 完全匹配检查
|
|
|
if any(col == field for field in row_fields for col in matched_columns):
|
|
|
return True
|
|
|
if any(col == field for field in row_fields for col in time_columns):
|
|
|
return True
|
|
|
if any(keyword in row for keyword in ["Code", "Abbr", "Name"]):
|
|
|
return True
|
|
|
return False
|
|
|
|
|
|
# 逐行过滤数据部分
|
|
|
filtered_rows = []
|
|
|
for row in data_rows.strip().split("\n"):
|
|
|
if satisfies_conditions(row):
|
|
|
filtered_rows.append(row)
|
|
|
|
|
|
# 将过滤后的内容与表头合并
|
|
|
filtered_structure = header + "\n\n" + "\n".join(filtered_rows) if filtered_rows else header
|
|
|
return filtered_structure
|
|
|
|
|
|
def extract_column_descriptions(filtered_structure, df_tmp):
|
|
|
# 从"\n\n"开始提取内容部分
|
|
|
content = filtered_structure.split("\n\n", 1)[1].strip()
|
|
|
|
|
|
# 提取列名（每行第一个空格前的部分）
|
|
|
column_names = []
|
|
|
for line in content.split("\n"):
|
|
|
column_name = line.split()[0] # 获取第一个空格前的内容
|
|
|
column_names.append(column_name)
|
|
|
|
|
|
# 转换 df_tmp 为字典形式，方便查找
|
|
|
column_dict = dict(zip(df_tmp["column_name"], df_tmp["注释"]))
|
|
|
|
|
|
# 构造结果列表
|
|
|
result = []
|
|
|
for column_name in column_names:
|
|
|
if column_name in column_dict and len(str(column_dict[column_name]))>3:
|
|
|
result.append({column_name: column_dict[column_name]})
|
|
|
|
|
|
return result
|
|
|
|
|
|
def calculate_similarity(text_list, local_vectors=False):
|
|
|
"""
|
|
|
批量计算相似度：
|
|
|
- 输入一个文本列表，其中第一个文本作为基准，其余文本与基准比较
|
|
|
- 当候选文本超过 64 条时，会分批请求，最后返回所有候选文本与基准文本的相似度（保留4位小数）
|
|
|
"""
|
|
|
|
|
|
base_text = text_list[0]
|
|
|
# 先单独请求基准文本的 embedding
|
|
|
base_response = client.embeddings.create(
|
|
|
model="embedding-3",
|
|
|
input=[base_text]
|
|
|
)
|
|
|
base_embedding = base_response.data[0].embedding
|
|
|
base_embedding = np.array(base_embedding)
|
|
|
|
|
|
all_similarities = []
|
|
|
if local_vectors:
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f'------>local_vectors:{local_vectors}')
|
|
|
# 使用本地保存的向量
|
|
|
with open(local_vectors, 'r', encoding='utf-8') as f:
|
|
|
local_embeddings = json.load(f)
|
|
|
|
|
|
# 确保读取到的向量与候选文本对应
|
|
|
candidate_embeddings = np.array(local_embeddings)
|
|
|
else:
|
|
|
# 批量处理候选文本，每次最多请求64条
|
|
|
candidate_texts = text_list[1:]
|
|
|
candidate_embeddings = []
|
|
|
chunk_size = 64 # 每次最多请求64条
|
|
|
for i in range(0, len(candidate_texts), chunk_size):
|
|
|
chunk = candidate_texts[i:i+chunk_size]
|
|
|
response = client.embeddings.create(
|
|
|
model="embedding-3",
|
|
|
input=chunk
|
|
|
)
|
|
|
# 提取候选文本的 embedding 并转换为 NumPy 数组
|
|
|
embeddings = [item.embedding for item in response.data]
|
|
|
candidate_embeddings.extend(embeddings)
|
|
|
|
|
|
candidate_embeddings = np.array(candidate_embeddings)
|
|
|
|
|
|
# 计算余弦相似度： dot / (||base|| * ||candidate||)
|
|
|
dot_products = candidate_embeddings.dot(base_embedding)
|
|
|
norm_base = np.linalg.norm(base_embedding)
|
|
|
norm_candidates = np.linalg.norm(candidate_embeddings, axis=1)
|
|
|
similarities = dot_products / (norm_base * norm_candidates)
|
|
|
|
|
|
# 保留4位小数，并加入结果列表
|
|
|
sims = [round(float(sim), 4) for sim in similarities]
|
|
|
all_similarities.extend(sims)
|
|
|
|
|
|
return all_similarities
|
