55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
|
import json
|
|||
|
|
import os
|
|||
|
|
import argparse
|
|||
|
|
import pandas as pd
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
|
|||
|
|
def get_Composition_ratio(input_file):
|
|||
|
|
"""
|
|||
|
|
计算数据集类别组成比例，并打印输出。
|
|||
|
|
:param input_file: 输入的JSONL文件路径
|
|||
|
|
"""
|
|||
|
|
# 读取JSONL文件
|
|||
|
|
with open(input_file, "r", encoding="utf-8") as f:
|
|||
|
|
data = [json.loads(line) for line in f]
|
|||
|
|
|
|||
|
|
# 提取每条数据的类别标签（假设在 conversation[0]['assistant']）
|
|||
|
|
labels = []
|
|||
|
|
for item in data:
|
|||
|
|
# 兼容 conversation 为列表且有 assistant 字段
|
|||
|
|
if "conversation" in item and isinstance(item["conversation"], list):
|
|||
|
|
conv = item["conversation"]
|
|||
|
|
if len(conv) > 0 and "assistant" in conv[0]:
|
|||
|
|
labels.append(conv[0]["assistant"])
|
|||
|
|
else:
|
|||
|
|
labels.append("未知")
|
|||
|
|
else:
|
|||
|
|
labels.append("未知")
|
|||
|
|
|
|||
|
|
df = pd.DataFrame({"label": labels})
|
|||
|
|
|
|||
|
|
# 计算每个类别的数量
|
|||
|
|
counts = df['label'].value_counts()
|
|||
|
|
total = counts.sum()
|
|||
|
|
|
|||
|
|
# 计算每个类别的比例
|
|||
|
|
ratios = counts / total * 100
|
|||
|
|
|
|||
|
|
# 打印每个类别的比例
|
|||
|
|
print("类别比例和数量:")
|
|||
|
|
for category, ratio in ratios.items():
|
|||
|
|
print(f"类别 {category}: {ratio:.2f}% ({counts[category]} 条)")
|
|||
|
|
|
|||
|
|
# 绘制饼图
|
|||
|
|
plt.figure(figsize=(8, 6))
|
|||
|
|
plt.pie(ratios, labels=ratios.index, autopct='%1.1f%%', startangle=140)
|
|||
|
|
plt.title('数据集类别比例')
|
|||
|
|
plt.show()
|
|||
|
|
return ratios
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# input_file = "sftdata.jsonl"
|
|||
|
|
input_file = "output-26.jsonl"
|
|||
|
|
input_file = "arxiv-metadata-oai-snapshot--swift-26.json"
|
|||
|
|
get_Composition_ratio(input_file)
|