71 lines
2.9 KiB
Python
71 lines
2.9 KiB
Python
|
import json
|
|||
|
|
from collections import defaultdict
|
|||
|
|
import pandas as pd
|
|||
|
|
|
|||
|
|
def classify_data_by_cluster_center(input_file, output_jsonl, output_excel):
|
|||
|
|
"""
|
|||
|
|
根据 cluster_center 对数据进行分类。
|
|||
|
|
每个 cluster_center 对应一个列表，存储去掉了 embedding 的数据。
|
|||
|
|
输出 JSONL 文件按类型的数据量从多到少排序，并记录每种类型的数量。
|
|||
|
|
同时将类型及其数量导出到 Excel 表格中。
|
|||
|
|
每个 cluster 对应的 data_list 会根据 'score' 字段降序排序。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 初始化分类字典
|
|||
|
|
classified_data = defaultdict(list)
|
|||
|
|
|
|||
|
|
# 读取输入文件
|
|||
|
|
with open(input_file, 'r', encoding='utf-8') as f:
|
|||
|
|
for line in f:
|
|||
|
|
record = json.loads(line.strip())
|
|||
|
|
|
|||
|
|
# 提取 cluster_center 部分
|
|||
|
|
cluster_center = record.get("cluster_center")
|
|||
|
|
|
|||
|
|
# 如果 cluster_center 存在，则根据其值分类
|
|||
|
|
if cluster_center is not None:
|
|||
|
|
record_without_embedding = {k: v for k, v in record.items() if k != "embedding"}
|
|||
|
|
classified_data[cluster_center].append(record_without_embedding)
|
|||
|
|
else:
|
|||
|
|
# 如果没有 cluster_center，则归类到 "null"
|
|||
|
|
record_without_embedding = {k: v for k, v in record.items() if k != "embedding"}
|
|||
|
|
classified_data["null"].append(record_without_embedding)
|
|||
|
|
|
|||
|
|
# 对每个 cluster_center 下的 data_list 按照 score 排序（默认为0）
|
|||
|
|
for center in classified_data:
|
|||
|
|
classified_data[center].sort(key=lambda x: x.get('score', 0), reverse=True)
|
|||
|
|
|
|||
|
|
# 按类型的数据量从多到少排序
|
|||
|
|
sorted_classified_data = sorted(classified_data.items(), key=lambda x: len(x[1]), reverse=True)
|
|||
|
|
|
|||
|
|
# 写入 JSONL 文件
|
|||
|
|
total_types = len(sorted_classified_data)
|
|||
|
|
with open(output_jsonl, 'w', encoding='utf-8') as out_f:
|
|||
|
|
for cluster_center, data_list in sorted_classified_data:
|
|||
|
|
entry = {
|
|||
|
|
str(cluster_center): data_list,
|
|||
|
|
#"count": len(data_list)
|
|||
|
|
}
|
|||
|
|
out_f.write(json.dumps(entry, ensure_ascii=False) + '\n')
|
|||
|
|
|
|||
|
|
# 准备 Excel 数据
|
|||
|
|
excel_data = []
|
|||
|
|
for cluster_center, data_list in sorted_classified_data:
|
|||
|
|
excel_data.append({"Cluster Center": cluster_center, "Count": len(data_list)})
|
|||
|
|
|
|||
|
|
# 导出到 Excel 文件
|
|||
|
|
df = pd.DataFrame(excel_data)
|
|||
|
|
df.to_excel(output_excel, index=False)
|
|||
|
|
|
|||
|
|
print(f"Total types: {total_types}")
|
|||
|
|
return total_types
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 示例用法
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
input_file = './dhbq/dhbq_merged_with_score_0513.jsonl'
|
|||
|
|
output_jsonl = './dhbq/dhbq_count_cluster_0513.jsonl'
|
|||
|
|
output_excel = './dhbq/dhbq_count_cluster_0513.xlsx'
|
|||
|
|
|
|||
|
|
total_types = classify_data_by_cluster_center(input_file, output_jsonl, output_excel)
|
|||
|
|
print(f"Total types found: {total_types}")
|