import sys import os sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import json from pathlib import Path from tqdm import tqdm import pandas as pd import numpy as np from utils.openai_access import call_chatgpt from utils.mpr import MultipleProcessRunnerSimplifier from utils.prompts import LLM_SCORE_PROMPT import re qa_data = {} def load_qa_results_from_dir(results_dir): """从结果目录加载所有QA结果""" global qa_data qa_data = {} results_path = Path(results_dir) json_files = list(results_path.glob("*.json")) print(f"找到 {len(json_files)} 个结果文件") for json_file in tqdm(json_files, desc="加载QA结果"): try: with open(json_file, 'r') as f: data = json.load(f) if ('index' in data and 'protein_id' in data and 'ground_truth' in data and 'llm_answer' in data): qa_data[data['index']] = data except Exception as e: print(f"加载文件 {json_file} 时出错: {e}") print(f"成功加载 {len(qa_data)} 个QA对") return qa_data def extract_score_from_response(response): """从LLM响应中提取分数""" if not response: return None # 尝试解析JSON格式的响应 try: if isinstance(response, str): # 尝试直接解析JSON json_match = re.search(r'\{[^}]*"score"[^}]*\}', response) if json_match: json_obj = json.loads(json_match.group()) return json_obj.get('score') # 尝试提取数字 score_match = re.search(r'"score":\s*(\d+(?:\.\d+)?)', response) if score_match: return float(score_match.group(1)) # 尝试提取纯数字 number_match = re.search(r'(\d+(?:\.\d+)?)', response) if number_match: score = float(number_match.group(1)) if 0 <= score <= 100: return score elif isinstance(response, dict): return response.get('score') except: pass return None def process_single_scoring(process_id, idx, qa_index, writer, save_dir): """处理单个QA对的打分""" try: qa_item = qa_data[qa_index] protein_id = qa_item['protein_id'] question = qa_item.get('question', '') ground_truth = qa_item['ground_truth'] llm_answer = qa_item['llm_answer'] # 构建打分提示 scoring_prompt = LLM_SCORE_PROMPT.replace('{{ground_truth}}', str(ground_truth)) scoring_prompt = scoring_prompt.replace('{{llm_answer}}', str(llm_answer)) # 调用LLM进行打分 score_response = call_chatgpt(scoring_prompt) score = extract_score_from_response(score_response) # 构建结果数据 result = { 'index': qa_index, 'protein_id': protein_id, 'question': question, 'ground_truth': ground_truth, 'llm_answer': llm_answer, 'score': score, 'raw_score_response': score_response } # 保存文件 save_path = os.path.join(save_dir, f"score_{protein_id}_{qa_index}.json") with open(save_path, 'w') as f: json.dump(result, f, indent=2, ensure_ascii=False) except Exception as e: print(f"处理QA索引 {qa_index} 时出错: {str(e)}") def get_missing_score_indices(save_dir): """检查哪些QA索引尚未完成打分""" all_qa_indices = list(qa_data.keys()) problem_qa_indices = set() for qa_index in tqdm(all_qa_indices, desc="检查打分文件"): protein_id = qa_data[qa_index]['protein_id'] json_file = Path(save_dir) / f"score_{protein_id}_{qa_index}.json" if not json_file.exists(): problem_qa_indices.add(qa_index) continue try: with open(json_file, 'r') as f: data = json.load(f) if (data is None or len(data) == 0 or 'score' not in data or data.get('score') is None): problem_qa_indices.add(qa_index) json_file.unlink() except Exception as e: problem_qa_indices.add(qa_index) try: json_file.unlink() except: pass return problem_qa_indices def collect_scores_to_json(save_dir, output_json): """收集所有打分结果并保存为JSON文件""" results = [] save_path = Path(save_dir) score_files = list(save_path.glob("score_*.json")) for score_file in tqdm(score_files, desc="收集打分结果"): try: with open(score_file, 'r') as f: data = json.load(f) results.append({ 'index': data.get('index'), 'protein_id': data.get('protein_id'), 'question': data.get('question', ''), 'ground_truth': data.get('ground_truth'), 'llm_answer': data.get('llm_answer'), 'score': data.get('score') }) except Exception as e: print(f"读取文件 {score_file} 时出错: {e}") # 按index排序 results.sort(key=lambda x: x.get('index', 0)) # 保存为JSON文件 with open(output_json, 'w', encoding='utf-8') as f: json.dump(results, f, indent=2, ensure_ascii=False) print(f"打分结果已保存到: {output_json}") # 转换为DataFrame用于分析 df = pd.DataFrame(results) return df def analyze_scores(df): """对打分结果进行统计分析""" print("\n=== 打分结果统计分析 ===") # 基本统计 valid_scores = df[df['score'].notna()]['score'] if len(valid_scores) == 0: print("没有有效的打分结果") return print(f"总样本数: {len(df)}") print(f"有效打分数: {len(valid_scores)}") print(f"无效打分数: {len(df) - len(valid_scores)}") print(f"有效率: {len(valid_scores)/len(df)*100:.2f}%") print(f"\n分数统计:") print(f"平均分: {valid_scores.mean():.2f}") print(f"中位数: {valid_scores.median():.2f}") print(f"标准差: {valid_scores.std():.2f}") print(f"最高分: {valid_scores.max():.2f}") print(f"最低分: {valid_scores.min():.2f}") # 分数分布 print(f"\n分数分布:") bins = [0, 20, 40, 60, 80, 100] labels = ['0-20', '21-40', '41-60', '61-80', '81-100'] for i, (low, high) in enumerate(zip(bins[:-1], bins[1:])): count = len(valid_scores[(valid_scores >= low) & (valid_scores <= high)]) percentage = count / len(valid_scores) * 100 print(f"{labels[i]}: {count} ({percentage:.1f}%)") # 分位数 print(f"\n分位数:") quantiles = [0.1, 0.25, 0.5, 0.75, 0.9] for q in quantiles: print(f"{int(q*100)}%分位数: {valid_scores.quantile(q):.2f}") # 按蛋白质ID分析（如果样本足够多） if len(df['protein_id'].unique()) > 1: print(f"\n按蛋白质ID分析:") protein_stats = df[df['score'].notna()].groupby('protein_id')['score'].agg(['count', 'mean', 'std']).round(2) print(protein_stats.head(10)) # 保存统计分析结果 stats_result = { "basic_stats": { "total_samples": len(df), "valid_scores": len(valid_scores), "invalid_scores": len(df) - len(valid_scores), "valid_rate": len(valid_scores)/len(df)*100, "mean_score": float(valid_scores.mean()), "median_score": float(valid_scores.median()), "std_score": float(valid_scores.std()), "max_score": float(valid_scores.max()), "min_score": float(valid_scores.min()) }, "distribution": {}, "quantiles": {} } # 分数分布统计 for i, (low, high) in enumerate(zip(bins[:-1], bins[1:])): count = len(valid_scores[(valid_scores >= low) & (valid_scores <= high)]) percentage = count / len(valid_scores) * 100 stats_result["distribution"][labels[i]] = { "count": count, "percentage": percentage } # 分位数统计 for q in quantiles: stats_result["quantiles"][f"{int(q*100)}%"] = float(valid_scores.quantile(q)) return stats_result def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument("--results_dir", type=str, default="data/evolla_hard_motif_go", help="包含LLM答案结果的目录") parser.add_argument("--n_process", type=int, default=32, help="并行进程数") parser.add_argument("--save_dir", type=str, default="data/llm_scores", help="保存打分结果的目录") parser.add_argument("--output_json", type=str, default="data/llm_scores_results.json", help="输出JSON文件路径") parser.add_argument("--stats_json", type=str, default="data/llm_scores_stats.json", help="统计分析结果JSON文件路径") parser.add_argument("--max_iterations", type=int, default=3, help="最大迭代次数") args = parser.parse_args() # 创建保存目录 os.makedirs(args.save_dir, exist_ok=True) os.makedirs(os.path.dirname(args.output_json), exist_ok=True) # 加载QA结果数据 load_qa_results_from_dir(args.results_dir) if not qa_data: print("没有找到有效的QA结果数据") return # 循环检查和打分 iteration = 0 while iteration < args.max_iterations: iteration += 1 print(f"\n开始第 {iteration} 轮打分") # 获取缺失打分的QA索引 missing_indices = get_missing_score_indices(args.save_dir) if not missing_indices: print("所有QA对已完成打分！") break print(f"发现 {len(missing_indices)} 个待打分的QA对") missing_indices_list = sorted(list(missing_indices)) # 使用多进程处理打分 mprs = MultipleProcessRunnerSimplifier( data=missing_indices_list, do=lambda process_id, idx, qa_index, writer: process_single_scoring(process_id, idx, qa_index, writer, args.save_dir), n_process=args.n_process, split_strategy="static" ) mprs.run() print(f"第 {iteration} 轮打分完成") # 收集结果并保存为JSON df = collect_scores_to_json(args.save_dir, args.output_json) # 进行统计分析 stats_result = analyze_scores(df) # 保存统计分析结果为JSON with open(args.stats_json, 'w', encoding='utf-8') as f: json.dump(stats_result, f, indent=2, ensure_ascii=False) print(f"统计分析结果已保存到: {args.stats_json}") # 检查最终结果 final_missing = get_missing_score_indices(args.save_dir) if final_missing: print(f"\n仍有 {len(final_missing)} 个QA对未能成功打分") else: print(f"\n所有 {len(qa_data)} 个QA对已成功完成打分！") if __name__ == "__main__": main()