format_data.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}

### 標準參考答案 (Ground Truth):
{reference_answer}

### 待評分的回覆:
{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):
        # 我們只處理 `traditional_*.json` 檔案
        if not filename.startswith("traditional_") or not filename.endswith(".json"):
            continue
        
        file_path = os.path.join(source_dir, filename)
        logging.info(f"  正在處理檔案: {filename}")
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)

            for item in data:
                mission_type = item.get('mission', 'unknown')
                original_prompt = item.get('prompt')
                reference_answer = item.get('resp')
                response_to_judge = item.get('model_response')
                ground_truth_judgement = item.get('full_output')

                if all([mission_type, original_prompt, reference_answer, response_to_judge, ground_truth_judgement]):
                    
                    # ---建立 user/assistant 格式 ---
                    
                    user_content = JUDGE_PROMPT_TEMPLATE.format(
                        original_prompt=original_prompt,
                        reference_answer=reference_answer,
                        response_to_be_judged=response_to_judge
                    ).strip() # 使用 strip() 移除前後多餘的空白換行
                    
                    # assistant_content 直接使用 ground_truth_judgement
                    assistant_content = ground_truth_judgement

                    # 3. 組成 LLaMA-Factory 偏好的 ShareGPT/messages 格式
                    formatted_item = {
                        "messages": [
                            {"role": "user", "content": user_content},
                            {"role": "assistant", "content": assistant_content}
                        ]
                    }
                    formatted_data.append(formatted_item)
                else:
                    logging.warning(f"    跳過一筆資料，因為缺少 'prompt', 'resp', 'model_response', 或 'full_output'。 (QID: {item.get('qid')})")
            
            total_processed_count += len(data)

        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 = "final_classified_output_new"
    OUTPUT_DIRECTORY = "judge_training_data"
    
    format_data_for_judge_model(
        source_dir=SOURCE_DIRECTORY,
        output_dir=OUTPUT_DIRECTORY,
        output_filename="judge_finetuning_dataset.jsonl"
    )