training_data_generator.py

import json
import random
import os
from typing import List, Dict, Any
import re

class HybridDataGenerator:
    def __init__(self, data_root: str = "../data"):
        self.think_start = "<|think|>"
        self.key_insight = "<|key-insight|>"
        self.think_continue = "<|think-continue|>"
        self.conclusion = "<|conclusion|>"
        self.data_root = data_root
    
    def load_gsm_data(self, split: str = "train") -> List[Dict[str, Any]]:
        file_path = os.path.join(self.data_root, f"gsm_{split}.json")
        
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Data file not found: {file_path}")
        
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        print(f"Loaded {len(data)} examples from {file_path}")
        return data
    
    def parse_calculation_steps(self, steps: List[str]) -> List[Dict[str, str]]:
        """解析计算步骤，提取计算和结果"""
        parsed_steps = []
        for step in steps:
            match = re.match(r'<<(.+?)=(.+?)>>', step)
            if match:
                calculation = match.group(1)
                result = match.group(2)
                parsed_steps.append({
                    'calculation': calculation,
                    'result': result,
                    'full_step': step
                })
        
        return parsed_steps
    
    def convert_to_hybrid_reasoning(self, example: Dict[str, Any]) -> str:
        """将GSM例子转换为混合推理格式"""
        question = example['question']
        steps = example['steps']
        answer = example['answer']
        parsed_steps = self.parse_calculation_steps(steps)
        
        if not parsed_steps:
            # 如果没有步骤，创建简单的推理
            reasoning = f"""
{self.think_start}
Let me analyze this problem step by step...
{self.key_insight} I need to work through the given information carefully.
{self.think_continue}
Processing the problem...
{self.conclusion} The answer is {answer}
"""
            return reasoning.strip()
        
        reasoning_parts = []
        reasoning_parts.append(self.think_start)
        reasoning_parts.append("Let me break down this problem step by step...")

        for i, step_info in enumerate(parsed_steps):
            if i == 0:
                reasoning_parts.append(f"{self.key_insight} First, I need to calculate: {step_info['calculation']}")
            elif i == len(parsed_steps) - 1:
                reasoning_parts.append(f"{self.think_continue}")
                reasoning_parts.append("Now for the final calculation...")
                reasoning_parts.append(f"{self.key_insight} Final step: {step_info['calculation']} = {step_info['result']}")
            else:
                reasoning_parts.append(f"{self.think_continue}")
                reasoning_parts.append("Moving to the next step...")
                reasoning_parts.append(f"{self.key_insight} Next: {step_info['calculation']} = {step_info['result']}")
        
        reasoning_parts.append(f"{self.conclusion} Therefore, the answer is {answer}")
        
        return "\n".join(reasoning_parts)
    
    def generate_hybrid_examples_from_gsm(self, split: str = "train", max_examples: int = None) -> List[Dict[str, Any]]:
        gsm_data = self.load_gsm_data(split)
        if max_examples is not None:
            gsm_data = gsm_data[:max_examples]
        
        hybrid_examples = []
        
        print(f"Converting {len(gsm_data)} GSM examples to hybrid format...")
        
        for i, example in enumerate(gsm_data):
            try:
                hybrid_reasoning = self.convert_to_hybrid_reasoning(example)
                hybrid_examples.append({
                    "question": example['question'],
                    "reasoning": hybrid_reasoning,
                    "answer": example['answer'],
                    "original_steps": example['steps'],
                    "type": "gsm_math"
                })
                if (i + 1) % 1000 == 0:
                    print(f"  Processed {i + 1}/{len(gsm_data)} examples...")
                    
            except Exception as e:
                print(f"Error processing example {i}: {e}")
                continue
        
        print(f"Generated {len(hybrid_examples)} hybrid examples from GSM data")
        return hybrid_examples
    
    def create_dataset_from_existing_data(self, save_path: str):
        print("Creating hybrid dataset from existing GSM8K data...")
        print("Loading GSM training data...")
        gsm_train_examples = self.generate_hybrid_examples_from_gsm("train")
        
        print("Loading GSM validation data...")
        gsm_val_examples = self.generate_hybrid_examples_from_gsm("valid")
        train_examples = gsm_train_examples
        val_examples = gsm_val_examples
        print("Formatting training data...")
        train_data = self.format_for_training(train_examples)
        print("Formatting validation data...")
        val_data = self.format_for_training(val_examples)
  
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        
        train_path = save_path.replace('.json', '_train.json')
        val_path = save_path.replace('.json', '_val.json')
        
        print("Saving training data...")
        with open(train_path, 'w', encoding='utf-8') as f:
            json.dump(train_data, f, indent=2, ensure_ascii=False)
        
        print("Saving validation data...")
        with open(val_path, 'w', encoding='utf-8') as f:
            json.dump(val_data, f, indent=2, ensure_ascii=False)
        
        print(f"Hybrid dataset created:")
        print(f"  Training: {len(train_data)} examples -> {train_path}")
        print(f"  Validation: {len(val_data)} examples -> {val_path}")
        
        return train_path, val_path
    
    def format_for_training(self, examples: List[Dict[str, Any]]) -> List[Dict[str, str]]:
        formatted = []
        for example in examples:
            text = f"Question: {example['question']}\n\nAnswer: {example['reasoning']}"
            
            formatted.append({
                "text": text,
                "question": example['question'],
                "answer": example['answer'],
                "type": example['type']
            })
        return formatted
    
    
    def create_dataset(self, save_path: str):
        return self.create_dataset_from_existing_data(save_path)


def extract_insights_from_reasoning(text: str) -> List[str]:
    """从推理过程中提取关键洞察"""
    pattern = r'<\|key-insight\|>\s*([^<]+)'
    insights = re.findall(pattern, text)
    
    return [insight.strip() for insight in insights]


def evaluate_reasoning_quality(example: Dict[str, Any]) -> Dict[str, float]:
    text = example.get('text', '')
    if 'Answer: ' in text:
        reasoning = text.split('Answer: ', 1)[1]
    else:
        reasoning = example.get('reasoning', '')
    num_insights = len(extract_insights_from_reasoning(reasoning))
    num_thinking_phases = reasoning.count('<|think')
    has_conclusion = '<|conclusion|>' in reasoning

    balance_score = min(num_insights / max(num_thinking_phases, 1), 1.0) if num_thinking_phases > 0 else 0.0
    
    return {
        'num_insights': num_insights,
        'num_thinking_phases': num_thinking_phases,
        'has_conclusion': has_conclusion,
        'balance_score': balance_score
    }


if __name__ == "__main__":
    generator = HybridDataGenerator(data_root="../data")
    
    train_path, val_path = generator.create_dataset_from_existing_data(
        save_path="./data/hybrid_dataset.json"
    )
    print("Done!")
    
    # with open(train_path, 'r', encoding='utf-8') as f:
    #     train_data = json.load(f)
    
    # print("\n" + "="*50)
    # print("示例训练数据:")
    # print("="*50)
    
    # for i, example in enumerate(train_data[:3]):
    #     print(f"\n示例 {i+1}:")
    #     print("-" * 30)
    #     print(example['text'])

    #     quality = evaluate_reasoning_quality(example)
    #     print(f"\n质量评估: {quality}")
    #     print("-" * 50)