PrepareInferenceData.py

import os
import json
import logging
from datetime import datetime

# --- 日誌設定 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', encoding='utf-8')

JUDGE_PROMPT_TEMPLATE = """
### 原始任務指令:
{original_prompt}

### 待評分的回覆:
{response_to_be_judged}
"""

def format_data_for_judge_model(source_dir, output_dir, output_filename="judge_training_dataset.jsonl"):
    """
    將清理好的資料轉換為「裁判模型」的 user/assistant 對話微調格式。
    """
    if not os.path.exists(source_dir):
        logging.error(f"錯誤：來源目錄 '{source_dir}' 不存在！")
        return

    os.makedirs(output_dir, exist_ok=True)
    
    formatted_data = []
    total_processed_count = 0

    logging.info("開始將資料格式化為「裁判模型」的 user/assistant 對話訓練集...")

    for filename in os.listdir(source_dir):
        # 只處理 .jsonl 檔案
        if not filename.endswith(".jsonl"):
            continue
        
        file_path = os.path.join(source_dir, filename)
        logging.info(f"  正在處理檔案: {filename}")
        
        # 每個檔案的計數器
        file_processed_count = 0
        max_per_file = 100  # 每個檔案最多處理100筆
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():  # 跳過空行
                        # 檢查這個檔案是否已經處理了100筆資料
                        if file_processed_count >= max_per_file:
                            logging.info(f"檔案 {filename} 已處理100筆資料，停止處理此檔案")
                            break
                            
                        item = json.loads(line)  # 每行是一個 JSON 物件
                        total_processed_count += 1
                        file_processed_count += 1
                        
                        mission_type = item.get('mission', 'unknown')
                        original_prompt = item.get('prompt')
                        reference_answer = item.get('resp')
                        response_to_judge = item.get('model_resp')

                        if all([mission_type, original_prompt, reference_answer, response_to_judge]):
                            
                            # ---建立 user/assistant 格式 ---
                            
                            user_content = JUDGE_PROMPT_TEMPLATE.format(
                                original_prompt=original_prompt,
                                response_to_be_judged=response_to_judge
                            ).strip() # 移除 reference_answer 參數，因為模板中沒有
                            

                            # 3. 組成 LLaMA-Factory 偏好的 ShareGPT/messages 格式
                            formatted_item = {
                                "messages": [
                                    {"role": "user", "content": user_content},
                                ]
                            }
                            formatted_data.append(formatted_item)
                        else:
                            logging.warning(f"    跳過一筆資料，因為缺少 'prompt', 'resp', 'model_resp'。 (QID: {item.get('qid')})")
            
            logging.info(f"檔案 {filename} 處理完成，共處理 {file_processed_count} 筆資料")

        except Exception as e:
            logging.error(f"處理檔案 {filename} 時發生錯誤: {e}")

    # 寫入最終的 .jsonl 檔案
    output_path = os.path.join(output_dir, output_filename)
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            for entry in formatted_data:
                f.write(json.dumps(entry, ensure_ascii=False) + '\n')
        
        logging.info("="*50)
        logging.info("JudgeModel訓練資料集準備完成！")
        logging.info(f"總共處理了 {total_processed_count} 筆來源資料。")
        logging.info(f"成功轉換並寫入了 {len(formatted_data)} 筆資料至:")
        logging.info(f"  --> {os.path.abspath(output_path)}")
        logging.info("="*50)

    except Exception as e:
        logging.error(f"寫入最終訓練檔案 {output_path} 時發生錯誤: {e}")


# --- 主程式執行區 (維持不變) ---
if __name__ == "__main__":
    SOURCE_DIRECTORY = "eval_results/gen_Gemma-3-27B-it/judge_llama-3.3-70B-Instruct"
    OUTPUT_DIRECTORY = "judge_training_data"
    
    format_data_for_judge_model(
        source_dir=SOURCE_DIRECTORY,
        output_dir=OUTPUT_DIRECTORY,
        output_filename="judge_inference_dataset.jsonl"
    )