283 lines
10 KiB
Python
283 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
统计octuple分词结果中每一列每个token的出现次数，并生成分析报告
|
||
"""
|
||
|
||
import os
|
||
import numpy as np
|
||
from pathlib import Path
|
||
from collections import defaultdict, Counter
|
||
from tqdm import tqdm
|
||
import json
|
||
|
||
# Octuple的列名定义
|
||
COLUMN_NAMES = [
|
||
"pitch", # 0: Pitch/PitchDrum
|
||
"position", # 1: Position
|
||
"bar", # 2: Bar
|
||
"velocity", # 3: Velocity
|
||
"duration", # 4: Duration
|
||
"program", # 5: Program
|
||
"tempo", # 6: Tempo
|
||
"timesig" # 7: TimeSignature
|
||
]
|
||
|
||
|
||
def load_octuple_data(data_dir):
|
||
"""
|
||
加载所有octuple分词后的.npz文件
|
||
|
||
Args:
|
||
data_dir: 数据目录路径，可以是单个目录或包含多个子目录的根目录
|
||
|
||
Returns:
|
||
list: 所有加载的numpy数组列表
|
||
"""
|
||
data_dir = Path(data_dir)
|
||
npz_files = []
|
||
|
||
# 如果目录存在，查找所有.npz文件
|
||
if data_dir.exists():
|
||
npz_files = list(data_dir.rglob("*.npz"))
|
||
|
||
if not npz_files:
|
||
print(f"警告: 在 {data_dir} 中未找到.npz文件")
|
||
return []
|
||
|
||
print(f"找到 {len(npz_files)} 个.npz文件，开始加载...")
|
||
|
||
all_data = []
|
||
for npz_file in tqdm(npz_files, desc="加载数据"):
|
||
try:
|
||
data = np.load(npz_file)['arr_0']
|
||
# 确保数据是二维数组 (num_tokens, num_columns)
|
||
if data.ndim == 2:
|
||
all_data.append(data)
|
||
elif data.ndim == 1:
|
||
# 如果是一维，可能需要reshape，但octuple应该是二维的
|
||
print(f"警告: {npz_file} 是一维数组，跳过")
|
||
except Exception as e:
|
||
print(f"错误: 加载 {npz_file} 时出错: {e}")
|
||
continue
|
||
|
||
return all_data
|
||
|
||
|
||
def count_tokens_by_column(all_data):
|
||
"""
|
||
统计每一列每个token的出现次数
|
||
|
||
Args:
|
||
all_data: 所有数据的列表，每个元素是一个numpy数组 (num_tokens, num_columns)
|
||
|
||
Returns:
|
||
dict: {column_index: Counter({token_value: count})}
|
||
"""
|
||
column_counters = defaultdict(Counter)
|
||
|
||
print("统计token出现次数...")
|
||
for data in tqdm(all_data, desc="处理数据"):
|
||
if data.size == 0:
|
||
continue
|
||
|
||
num_columns = data.shape[1] if data.ndim == 2 else 1
|
||
|
||
for col_idx in range(num_columns):
|
||
if data.ndim == 2:
|
||
tokens = data[:, col_idx]
|
||
else:
|
||
tokens = data
|
||
|
||
# 统计该列中每个token的出现次数
|
||
unique, counts = np.unique(tokens, return_counts=True)
|
||
for token, count in zip(unique, counts):
|
||
column_counters[col_idx][int(token)] += int(count)
|
||
|
||
return dict(column_counters)
|
||
|
||
|
||
def generate_report(column_counters, output_file=None):
|
||
"""
|
||
生成分析报告
|
||
|
||
Args:
|
||
column_counters: 统计结果字典
|
||
output_file: 输出文件路径（可选）
|
||
"""
|
||
report_lines = []
|
||
report_lines.append("=" * 80)
|
||
report_lines.append("OCTUPLE分词结果统计分析报告")
|
||
report_lines.append("=" * 80)
|
||
report_lines.append("")
|
||
|
||
# 总体统计
|
||
total_tokens = sum(sum(counter.values()) for counter in column_counters.values())
|
||
report_lines.append(f"总token数: {total_tokens:,}")
|
||
report_lines.append(f"分析的列数: {len(column_counters)}")
|
||
report_lines.append("")
|
||
|
||
# 每一列的详细统计
|
||
for col_idx in sorted(column_counters.keys()):
|
||
counter = column_counters[col_idx]
|
||
col_name = COLUMN_NAMES[col_idx] if col_idx < len(COLUMN_NAMES) else f"column_{col_idx}"
|
||
|
||
report_lines.append("-" * 80)
|
||
report_lines.append(f"列 {col_idx}: {col_name}")
|
||
report_lines.append("-" * 80)
|
||
|
||
total_in_column = sum(counter.values())
|
||
unique_tokens = len(counter)
|
||
min_token = min(counter.keys()) if counter else 0
|
||
max_token = max(counter.keys()) if counter else 0
|
||
|
||
report_lines.append(f" 总token数: {total_in_column:,}")
|
||
report_lines.append(f" 唯一token数: {unique_tokens:,}")
|
||
report_lines.append(f" Token值范围: [{min_token}, {max_token}]")
|
||
report_lines.append(f" 平均出现次数: {total_in_column / unique_tokens:.2f}" if unique_tokens > 0 else " 平均出现次数: N/A")
|
||
report_lines.append("")
|
||
|
||
# Top 20 最常见的token
|
||
report_lines.append(f" Top 20 最常见的token:")
|
||
top_tokens = counter.most_common(20)
|
||
for rank, (token, count) in enumerate(top_tokens, 1):
|
||
percentage = (count / total_in_column * 100) if total_in_column > 0 else 0
|
||
report_lines.append(f" {rank:2d}. Token {token:6d}: {count:10,} 次 ({percentage:5.2f}%)")
|
||
report_lines.append("")
|
||
|
||
# Top 20 最不常见的token（出现次数>0的）
|
||
report_lines.append(f" Top 20 最不常见的token (出现次数>0):")
|
||
bottom_tokens = counter.most_common()[-20:]
|
||
bottom_tokens.reverse()
|
||
for rank, (token, count) in enumerate(bottom_tokens, 1):
|
||
percentage = (count / total_in_column * 100) if total_in_column > 0 else 0
|
||
report_lines.append(f" {rank:2d}. Token {token:6d}: {count:10,} 次 ({percentage:5.2f}%)")
|
||
report_lines.append("")
|
||
|
||
# 分布统计
|
||
counts_list = list(counter.values())
|
||
if counts_list:
|
||
report_lines.append(f" 分布统计:")
|
||
report_lines.append(f" 最小出现次数: {min(counts_list):,}")
|
||
report_lines.append(f" 最大出现次数: {max(counts_list):,}")
|
||
report_lines.append(f" 中位数出现次数: {np.median(counts_list):,.0f}")
|
||
report_lines.append(f" 标准差: {np.std(counts_list):,.2f}")
|
||
report_lines.append("")
|
||
|
||
# 跨列分析
|
||
report_lines.append("=" * 80)
|
||
report_lines.append("跨列分析")
|
||
report_lines.append("=" * 80)
|
||
report_lines.append("")
|
||
|
||
for col_idx in sorted(column_counters.keys()):
|
||
counter = column_counters[col_idx]
|
||
col_name = COLUMN_NAMES[col_idx] if col_idx < len(COLUMN_NAMES) else f"column_{col_idx}"
|
||
total_in_column = sum(counter.values())
|
||
percentage = (total_in_column / total_tokens * 100) if total_tokens > 0 else 0
|
||
report_lines.append(f" {col_name:12s}: {total_in_column:12,} tokens ({percentage:5.2f}%)")
|
||
|
||
report_lines.append("")
|
||
report_lines.append("=" * 80)
|
||
report_lines.append("报告生成完成")
|
||
report_lines.append("=" * 80)
|
||
|
||
# 输出报告
|
||
report_text = "\n".join(report_lines)
|
||
print("\n" + report_text)
|
||
|
||
# 保存到文件
|
||
if output_file:
|
||
output_path = Path(output_file)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
f.write(report_text)
|
||
print(f"\n报告已保存到: {output_path}")
|
||
|
||
# 同时保存JSON格式的详细数据
|
||
if output_file:
|
||
json_output = output_path.with_suffix('.json')
|
||
json_data = {
|
||
'summary': {
|
||
'total_tokens': total_tokens,
|
||
'num_columns': len(column_counters)
|
||
},
|
||
'columns': {}
|
||
}
|
||
|
||
for col_idx in sorted(column_counters.keys()):
|
||
counter = column_counters[col_idx]
|
||
col_name = COLUMN_NAMES[col_idx] if col_idx < len(COLUMN_NAMES) else f"column_{col_idx}"
|
||
json_data['columns'][col_name] = {
|
||
'total_tokens': sum(counter.values()),
|
||
'unique_tokens': len(counter),
|
||
'token_counts': dict(counter),
|
||
'top_20': dict(counter.most_common(20)),
|
||
'bottom_20': dict(counter.most_common()[-20:])
|
||
}
|
||
|
||
with open(json_output, 'w', encoding='utf-8') as f:
|
||
json.dump(json_data, f, indent=2, ensure_ascii=False)
|
||
print(f"详细数据已保存到: {json_output}")
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
# 默认数据目录 - 可以根据需要修改
|
||
default_data_dir = "dataset/represented_data/tuneidx/tuneidx_msmidi"
|
||
|
||
# 可以指定具体的数据目录，例如：
|
||
data_dir = "dataset/represented_data/tuneidx/tuneidx_msmidi/oct8_resampled_v2"
|
||
# 或者使用默认目录扫描所有oct8目录
|
||
|
||
# import sys
|
||
# if len(sys.argv) > 1:
|
||
# data_dir = sys.argv[1]
|
||
# else:
|
||
# # 自动查找所有oct8目录
|
||
# base_dir = Path(default_data_dir)
|
||
# oct8_dirs = list(base_dir.rglob("oct8"))
|
||
# if oct8_dirs:
|
||
# print(f"找到以下oct8目录:")
|
||
# for i, d in enumerate(oct8_dirs, 1):
|
||
# print(f" {i}. {d}")
|
||
# if len(oct8_dirs) == 1:
|
||
# data_dir = str(oct8_dirs[0])
|
||
# print(f"\n使用目录: {data_dir}")
|
||
# else:
|
||
# # 使用第一个找到的目录，或者合并所有目录
|
||
# print(f"\n使用第一个目录: {oct8_dirs[0]}")
|
||
# print("如需分析其他目录，请指定路径作为参数")
|
||
# data_dir = str(oct8_dirs[0])
|
||
# else:
|
||
# data_dir = default_data_dir
|
||
# print(f"未找到oct8目录，使用默认目录: {data_dir}")
|
||
|
||
# 加载数据
|
||
all_data = load_octuple_data(data_dir)
|
||
|
||
if not all_data:
|
||
print("错误: 未加载到任何数据")
|
||
return
|
||
|
||
# 检查数据格式
|
||
if all_data:
|
||
sample = all_data[0]
|
||
print(f"\n数据格式检查:")
|
||
print(f" 样本形状: {sample.shape}")
|
||
print(f" 样本数据类型: {sample.dtype}")
|
||
print(f" 列数: {sample.shape[1] if sample.ndim == 2 else 1}")
|
||
print()
|
||
|
||
# 统计token出现次数
|
||
column_counters = count_tokens_by_column(all_data)
|
||
|
||
# 生成报告
|
||
output_file = "octuple_token_analysis_report_part.txt"
|
||
generate_report(column_counters, output_file)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|