src/lmr/generation/generation.py

#!/usr/bin/env python3
"""
Unified Generator with TinyGSM/GSM8K Support
完全按照成功版本的逻辑 - 使用 simple_math_problem 函数名
包含实时准确率显示功能
新增：5-fold 评估支持（config.five_fold = True）
"""

import re
import torch
import ast
import multiprocessing as mp
import time
from pathlib import Path
from tqdm import tqdm
from typing import List, Dict, Optional, Union

# Default prompts
DEFAULT_CHAT_PROMPT = """[USER] Hello, how are you today? [SYSTEM] I am"""

def _exec_worker(code: str, out_q: mp.Queue):
    """
    子进程执行 worker：
    - 在受限的 globals 下 exec(code)
    - 将 ('ok', result) 或 ('err', error_message) 放入 out_q
    """
    safe_globals = {
        "__builtins__": {
            "abs": abs, "round": round, "min": min, "max": max, "sum": sum,
            "int": int, "float": float, "str": str, "len": len, "range": range,
            "list": list, "dict": dict, "set": set, "tuple": tuple,
            "True": True, "False": False, "None": None
        }
    }
    safe_locals = {}
    try:
        exec(code, safe_globals, safe_locals)
        result = safe_locals.get('result', None)
        out_q.put(('ok', result))
    except Exception as e:
        out_q.put(('err', repr(e)))

class Generator:
    """
    统一生成器，支持：
    - chat: 标准对话生成
    - tinygsm dataset: 自动在 GSM8K test 上评估（执行生成的 Python 代码）
    支持五折评估：将数据拆成 5 份，分别作为 test fold 执行评估并保存每个 fold 的结果
    """
    
    def __init__(self, config, model, tokenizer, device="cuda", output_dir=None):
        self.config = config.generation
        self.model = model.eval().to(device)
        self.tokenizer = tokenizer
        self.device = device
        
        # 将 output_dir 保存为实例属性，供保存结果使用
        self.output_dir = Path(output_dir or getattr(self.config, 'output_dir', 'output'))
        
        # 执行子进程超时（秒），可在 config 中设置 exec_timeout
        self.exec_timeout = getattr(self.config, 'exec_timeout', 20)
        
        # 检测 dataset 类型
        self.dataset_name = getattr(config, 'dataset', None)
        if hasattr(self.dataset_name, 'dataset_name'):
            self.dataset_name = self.dataset_name.dataset_name
        
        self.is_tinygsm = self.dataset_name == "tinygsm"
        
        print(f"[Generator] Initialized")
        print(f"  Dataset: {self.dataset_name}")
        print(f"  Device: {device}")
        print(f"  Exec timeout: {self.exec_timeout}s")
        if self.is_tinygsm:
            print(f"  Mode: TinyGSM - Will evaluate on GSM8K test set with CODE EXECUTION")
            print(f"  Five-fold enabled: {getattr(self.config, 'five_fold', True)}")

    def generate(self, prompt: Optional[Union[str, List[str]]] = None):
        """
        主生成入口
        
        - 如果是 tinygsm dataset: 自动加载 GSM8K test 并评估
        - 否则: 标准 chat 生成
        """
        if self.is_tinygsm:
            return self._generate_tinygsm()
        else:
            return self._generate_chat(prompt)
    
    def _generate_chat(self, prompt: Optional[str] = None):
        """标准 chat 模式生成"""
        
        if prompt is None:
            prompt_text = getattr(self.config, "prompt", DEFAULT_CHAT_PROMPT)
        else:
            prompt_text = prompt
            
        max_len = getattr(self.config, "max_new_tokens", 512)
        temperature = getattr(self.config, "temperature", 0.7)
        top_p = getattr(self.config, "top_p", 0.9)
        return_gen_only = getattr(self.config, "return_generation_only", True)

        input_ids = torch.tensor(
            [self.tokenizer.encode(prompt_text)],
            dtype=torch.long,
            device=self.device
        )

        with torch.no_grad():
            outputs = self.model.generate(
                input_ids,
                max_generation_length=max_len,
                tokenizer=self.tokenizer,
                temperature=temperature,
                top_p=top_p,
                return_generation_only=return_gen_only
            )

        generated_texts = []
        for seq in outputs:
            text = self.tokenizer.decode(seq.tolist())
            generated_texts.append(text)

        if getattr(self.config, "save_to_file", False):
            self._save_generations(generated_texts, "generations.txt")

        for text in generated_texts:
            print(f"{prompt_text} [{text}]")
        
        return generated_texts
    
    def _generate_tinygsm(self):
        """TinyGSM mode: 在 GSM8K test set 上生成并验证"""
        print("\n" + "="*60)
        print("TinyGSM Dataset Detected")
        print("Evaluating on GSM8K Test Set")
        print("="*60 + "\n")
        
        print("Step 1: Loading GSM8K Test Set...")
        gsm8k_data = self._load_gsm8k_test()
        
        if not gsm8k_data or not gsm8k_data['questions']:
            print("❌ Failed to load GSM8K test set")
            return None
        
        questions = gsm8k_data['questions']
        ground_truth = gsm8k_data['answers']
        
        max_samples = getattr(self.config, 'max_samples', None)
        if max_samples is not None and len(questions) > max_samples:
            print(f"Limiting to {max_samples} samples")
            questions = questions[:max_samples]
            ground_truth = ground_truth[:max_samples]
        
        print(f"✓ Loaded {len(questions)} questions\n")
        
        # 如果配置了 five_fold，则走五折流程
        if getattr(self.config, 'five_fold', False):
            return self._generate_tinygsm_five_fold(questions, ground_truth)
        
        print("Step 2: Generating Solutions...")
        # 单次评估
        results = self._generate_solutions(questions, ground_truth)
        
        print("\nStep 3: Verifying Answers...")
        verification = self._verify_answers(results, ground_truth)
        results['verification'] = verification
        results['ground_truth'] = ground_truth
        
        if getattr(self.config, 'save_to_file', False):
            self._save_tinygsm_results(results)
        
        self._print_verification_summary(verification)
        
        return results

    def _load_gsm8k_test(self) -> Dict[str, List]:
        """加载 GSM8K test set"""
        try:
            from datasets import load_dataset
            
            print("Loading from HuggingFace...")
            ds = load_dataset("openai/gsm8k", "main", split="test")
            
            questions = []
            answers = []
            
            with tqdm(total=len(ds), desc="Loading", unit="sample") as pbar:
                for sample in ds:
                    q = (sample.get("question", "") or "").strip()
                    a = (sample.get("answer", "") or "").strip()
                    
                    if not q or not a:
                        pbar.update(1)
                        continue
                    
                    answer_num = self._extract_gsm8k_answer(a)
                    
                    questions.append(q)
                    answers.append(answer_num)
                    
                    pbar.update(1)
            
            return {'questions': questions, 'answers': answers}
            
        except Exception as e:
            print(f"❌ Error loading GSM8K: {e}")
            return {'questions': [], 'answers': []}

    def _extract_gsm8k_answer(self, answer_text: str) -> str:
        """从 GSM8K 答案中提取数值"""
        match = re.search(r'####\s*([+-]?[\d,]+\.?\d*)', answer_text)
        if match:
            return match.group(1).replace(',', '').strip()
        
        match = re.search(r'result\s*=\s*([+-]?[\d,]+\.?\d*)', answer_text, re.IGNORECASE)
        if match:
            return match.group(1).replace(',', '').strip()
        
        numbers = re.findall(r'([+-]?[\d,]+\.?\d*)', answer_text)
        if numbers:
            return numbers[-1].replace(',', '').strip()
        
        return "0"

    def _create_math_prompt(self, question: str) -> str:
        """创建结构化的数学问题提示"""
        return f"""<|bos|>Question: {question}

Solution:
def simple_math_problem() -> int:
    # {question}
    """

    def _generate_solutions(self, questions: List[str], ground_truth: List[str]) -> Dict:
        """
        为每个问题生成代码解决方案
        在循环中加入实时准确率计算和进度条更新
        """
        max_len = 512
        temperature = getattr(self.config, 'temperature', 0.7)
        top_p = getattr(self.config, 'top_p', 0.9)
        
        results = {
            'questions': [],
            'prompts': [],
            'generated_solutions': [],
            'extracted_code': [],
            'executed_answers': [],
            'full_outputs': []
        }
        
        correct_count = 0
        total_processed = 0
        
        # 将问题和标准答案打包，并在 tqdm 中遍历
        pbar = tqdm(zip(questions, ground_truth), total=len(questions), desc="Generating", unit="q")
        
        for idx, (question, truth_val) in enumerate(pbar):
            print('before1')
            prompt = self._create_math_prompt(question)
            
            input_ids = torch.tensor(
                [self.tokenizer.encode(prompt)],
                dtype=torch.long,
                device=self.device
            )
            print('before')
            # import pdb
            # pdb.set_trace()
            if(input_ids is tuple):
                input_ids = input_ids.logits
            # import pdb
            # pdb.set_trace()
            with torch.no_grad():
                # try:
                    outputs = self.model.generate(
                        input_ids,
                        max_generation_length=max_len,
                        tokenizer=self.tokenizer,
                        temperature=temperature,
                        top_p=top_p,
                        return_generation_only=False
                    )
                # except:
                #     outputs = self.model.generate(
                #     input_ids,
                #     tokenizer=self.tokenizer,
                #     temperature=temperature,
                #     top_p=top_p,
                # )
            print('after')
            full_text = self.tokenizer.decode(outputs[0].tolist())
            print('after1')
            # 调试打印 (减少频率)
            if idx == 0:
                print(f"\n{'='*60}")
                print(f"🔍 DEBUG Sample {idx + 1}: Raw generated text (first 500 chars):")
                print(f"{'='*60}")
                print(full_text[:500])
                print(f"{'='*60}\n")
            
            if full_text.startswith(prompt):
                solution = full_text[len(prompt):]
            else:
                solution = full_text
            
            solution = solution.replace("<|eos|>", "").replace("<|bos|>", "").strip()
            
            code = self._extract_function_code(full_text)
            print('after2')
            
            if code:
                answer = self._execute_math_code(code)
            else:
                answer = None
            print('after3')
            # --- 实时验证和进度条更新 ---
            is_correct = False
            if truth_val is not None:
                is_correct = self._compare_answers(answer, truth_val)
                if is_correct:
                    correct_count += 1
                total_processed += 1
                
                # 更新进度条后缀显示准确率
                current_acc = correct_count / total_processed
                pbar.set_postfix({"Acc": f"{current_acc:.2%}"})
            # ------------------------------
            print('after4')
            results['questions'].append(question)
            results['prompts'].append(prompt)
            results['generated_solutions'].append(solution)
            results['extracted_code'].append(code if code else "No code extracted")
            results['executed_answers'].append(answer)
            results['full_outputs'].append({
                'question': question,
                'solution': solution,
                'code': code,
                'answer': answer
            })
            
            if idx < 3 or idx == len(questions) - 1:
                print(f"\n{'─'*60}")
                print(f"Sample {idx + 1}")
                print(f"{'─'*60}")
                print(f"Q: {question[:100]}...")
                if code:
                    print(f"Code:\n{code[:200]}...")
                print(f"Answer: {answer if answer is not None else '[None]'}")
                if truth_val:
                    status = "✓" if is_correct else "✗"
                    print(f"Truth:  {truth_val} {status}")
                print(f"{'─'*60}")
        
        return results

    # -------------------- 新增/修改方法：五折支持 --------------------
    def _split_into_folds(self, questions: List[str], answers: List[str], n_folds: int = 5):
        """按顺序把 questions/answers 拆成 n_folds 份，尽量平均"""
        assert len(questions) == len(answers)
        total = len(questions)
        folds_q = []
        folds_a = []
        base = total // n_folds
        remainder = total % n_folds
        start = 0
        for i in range(n_folds):
            size = base + (1 if i < remainder else 0)
            end = start + size
            folds_q.append(questions[start:end])
            folds_a.append(answers[start:end])
            start = end
        return folds_q, folds_a

    def _generate_tinygsm_five_fold(self, questions: List[str], ground_truth: List[str], n_folds: int = 5):
        """
        五折评估：
        - 将数据分成 n_folds 份
        - 对每个 fold 单独作为测试集，调用 _generate_solutions
        - 保存每个 fold 的结果到 self.output_dir/gsm8k_results_fold_{i+1}.txt
        - 返回包含每个 fold 详细结果和总体汇总的字典
        """
        print("Running 5-fold evaluation...")
        folds_q, folds_a = self._split_into_folds(questions, ground_truth, n_folds=n_folds)
        
        all_folds_results = []
        fold_accuracies = []
        
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        for i in range(n_folds):
            # if (i!=4):
                # continue
            print(f"\n{'#'*40}\nEvaluating Fold {i+1}/{n_folds} - test size: {len(folds_q[i])}\n{'#'*40}\n")
            
            # 对当前 fold 执行生成与验证
            # print
            fold_results = self._generate_solutions(folds_q[i], folds_a[i])
            
            # 验证（_verify_answers 返回的 total/accuracy 等）
            fold_verification = self._verify_answers(fold_results, folds_a[i])
            fold_results['verification'] = fold_verification
            fold_results['ground_truth'] = folds_a[i]
            
            # 保存 fold 文件
            fold_out_path = self.output_dir / f"gsm8k_results_fold_{i+1}.txt"
            with open(fold_out_path, 'w', encoding='utf-8') as f:
                for j, item in enumerate(fold_results['full_outputs'], 1):
                    f.write(f"{'='*60}\n")
                    f.write(f"Fold {i+1} - Problem {j}\n")
                    f.write(f"{'='*60}\n")
                    f.write(f"Question:\n{item['question']}\n\n")
                    f.write(f"Solution:\n{item['solution']}\n\n")
                    if item.get('code'):
                        f.write(f"Extracted Code:\n{item['code']}\n\n")
                    f.write(f"Answer: {item['answer']}\n")
                    if 'ground_truth' in fold_results and j <= len(fold_results['ground_truth']):
                        truth = fold_results['ground_truth'][j-1]
                        match = "✓" if self._compare_answers(item['answer'], truth) else "✗"
                        f.write(f"Truth: {truth} {match}\n")
                    f.write(f"{'='*60}\n\n")
            print(f"[Generator] Saved fold results to {fold_out_path}")

            # 保存 verification 简要文件
            verify_path = self.output_dir / f"verification_fold_{i+1}.txt"
            with open(verify_path, 'w', encoding='utf-8') as f:
                v = fold_verification
                f.write(f"{'='*60}\n")
                f.write(f"GSM8K Verification Results - Fold {i+1}\n")
                f.write(f"{'='*60}\n")
                f.write(f"Total: {v['total']}\n")
                f.write(f"Correct: {v['correct']}\n")
                f.write(f"Incorrect: {v['incorrect']}\n")
                f.write(f"No Answer: {v['no_answer']}\n")
                f.write(f"Accuracy: {v['accuracy']*100:.2f}%\n")
                f.write(f"{'='*60}\n\n")
            print(f"[Generator] Saved fold verification to {verify_path}")
            
            all_folds_results.append(fold_results)
            fold_accuracies.append(fold_verification.get('accuracy', 0.0))
            
            # 打印单折汇总
            print(f"Fold {i+1} Accuracy: {fold_verification.get('accuracy', 0.0)*100:.2f}%")
        
        # 总体汇总
        overall_mean_acc = sum(fold_accuracies) / len(fold_accuracies) if fold_accuracies else 0.0
        summary = {
            'n_folds': n_folds,
            'fold_accuracies': fold_accuracies,
            'mean_accuracy': overall_mean_acc,
            'folds': all_folds_results
        }
        
        # 保存总体 summary
        summary_path = self.output_dir / "gsm8k_5fold_summary.txt"
        with open(summary_path, 'w', encoding='utf-8') as f:
            f.write(f"{'='*60}\n")
            f.write("GSM8K 5-Fold Summary\n")
            f.write(f"{'='*60}\n")
            for idx, acc in enumerate(fold_accuracies, 1):
                f.write(f"Fold {idx} Accuracy: {acc*100:.2f}%\n")
            f.write(f"\nMean Accuracy: {overall_mean_acc*100:.2f}%\n")
            f.write(f"{'='*60}\n")
        print(f"[Generator] Saved 5-fold summary to {summary_path}")
        
        # 打印总体结果
        print("\n" + "="*60)
        print("5-Fold Evaluation Summary")
        print(f"Mean Accuracy: {overall_mean_acc*100:.2f}%")
        for idx, acc in enumerate(fold_accuracies, 1):
            print(f"  Fold {idx}: {acc*100:.2f}%")
        print("="*60 + "\n")
        
        return summary
    # -------------------- 五折支持结束 --------------------

    def _extract_function_code(self, generated_text):
        """Extract and clean function code from generated text"""
        func_start = generated_text.find("def simple_math_problem")
        if func_start == -1:
            return None
            
        func_text = generated_text[func_start:]
        
        lines = func_text.split('\n')
        cleaned_lines = []
        
        for i, line in enumerate(lines):
            if i == 0:
                cleaned_lines.append("def simple_math_problem() -> int:")
            else:
                if not cleaned_lines and not line.strip():
                    continue
                    
                # Stop at next function or end markers
                stripped_line = line.strip()
                if stripped_line and not line.startswith(' ') and not line.startswith('\t') and i > 0:
                    if 'def ' in line or stripped_line == 'result = simple_math_problem()':
                        break
                
                # Stop at markdown code block markers or answer markers
                if stripped_line in ['```', '```python', '```py'] or stripped_line.startswith('Answer:') or '<|eos|>' in stripped_line:
                    break
                
                clean_line = line.rstrip()
                original_line = clean_line
                
                if clean_line.strip():
                    # Check for end markers (even if indented)
                    stripped = clean_line.strip()
                    if stripped in ['```', '```python', '```py'] or stripped.startswith('Answer:') or '<|eos|>' in stripped:
                        break
                    
                    # skip docstrings and comments
                    if stripped.startswith('"""') or stripped.startswith("'''") or stripped.startswith('#'):
                        continue
                    
                    if clean_line.strip() and not clean_line.startswith('    '):
                        clean_line = '    ' + clean_line.strip()
                    
                    cleaned_lines.append(clean_line)
                else:
                    cleaned_lines.append('')
        
        while cleaned_lines and not cleaned_lines[-1].strip():
            cleaned_lines.pop()
            
        has_return = any('return' in line for line in cleaned_lines)
        if not has_return:
            cleaned_lines.append('    return result')
        
        final_code = '\n'.join(cleaned_lines)
        return final_code

    def _execute_math_code(self, code):
        """Safely execute mathematical code with a timeout using multiprocessing."""
        try:
            if not code:
                return None

            # 确保函数体有 return（原逻辑）
            if 'return' not in code:
                code += '\n    return result'

            full_code = code + '\n\nresult = simple_math_problem()'

            # 先简单做语法检查
            try:
                ast.parse(full_code)
            except SyntaxError:
                return None

            # 创建 Queue 与 Process
            q: mp.Queue = mp.Queue()
            p = mp.Process(target=_exec_worker, args=(full_code, q))
            p.start()

            # 等待，超时则 terminate
            p.join(timeout=self.exec_timeout)
            if p.is_alive():
                try:
                    p.terminate()
                    p.join(1)
                except Exception:
                    pass
                # 尝试清空队列
                try:
                    while not q.empty():
                        q.get_nowait()
                except Exception:
                    pass
                print(f"[_execute_math_code] Execution timed out after {self.exec_timeout}s.")
                return None

            # 子进程已结束，从队列读结果
            try:
                status, payload = q.get(timeout=0.1)
            except Exception:
                return None

            if status == 'ok':
                try:
                    result_value = float(payload) if payload is not None else None
                    return result_value
                except Exception:
                    return None
            else:
                print(f"[_execute_math_code] Child process error: {payload}")
                return None

        except Exception as e:
            print(f"[_execute_math_code] Unexpected error: {e!r}")
            return None

    def _verify_answers(self, results: Dict, ground_truth: List[str]) -> Dict:
        """验证答案"""
        predictions = results['executed_answers']
        
        verification = {
            'total': len(ground_truth),
            'correct': 0,
            'incorrect': 0,
            'no_answer': 0,
            'accuracy': 0.0,
            'details': []
        }
        
        for idx, (pred, truth) in enumerate(zip(predictions, ground_truth)):
            if pred is None:
                verification['no_answer'] += 1
                is_correct = False
            else:
                is_correct = self._compare_answers(pred, truth)
                
                if is_correct:
                    verification['correct'] += 1
                else:
                    verification['incorrect'] += 1
            
            verification['details'].append({
                'index': idx,
                'question': results['questions'][idx],
                'predicted': pred,
                'ground_truth': truth,
                'correct': is_correct
            })
        
        verification['accuracy'] = (
            verification['correct'] / verification['total']
            if verification['total'] > 0 else 0.0
        )
        
        return verification

    def _compare_answers(self, pred: str, truth: str) -> bool:
        """数值比较"""
        pred_str = str(pred).strip().replace(',', '')
        truth_str = str(truth).strip().replace(',', '')
        
        if pred_str == truth_str:
            return True
        
        try:
            pred_num = float(pred_str)
            truth_num = float(truth_str)
            return abs(pred_num - truth_num) < 1e-6
        except (ValueError, TypeError):
            pass
        
        return False

    def _print_verification_summary(self, verification: Dict):
        """打印验证总结"""
        print(f"\n{'='*60}")
        print("GSM8K Evaluation Results")
        print(f"{'='*60}")
        print(f"Total:      {verification['total']}")
        print(f"Correct:    {verification['correct']} ({verification['accuracy']*100:.2f}%)")
        print(f"Incorrect:  {verification['incorrect']}")
        print(f"No Answer:  {verification['no_answer']}")
        print(f"{'='*60}")
        
        incorrect = [d for d in verification['details'] if not d['correct']][:3]
        if incorrect:
            print("\nSample Incorrect Predictions:")
            for detail in incorrect:
                print(f"  Q: {detail['question'][:60]}...")
                print(f"  Predicted: {detail['predicted']}, Truth: {detail['ground_truth']}\n")

    def _save_generations(self, texts: List[str], filename: str):
        """保存 chat 生成结果"""
        output_dir = Path(getattr(self.config, 'output_dir', self.output_dir))
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir / filename
        
        with open(output_path, 'w', encoding='utf-8') as f:
            for i, text in enumerate(texts, 1):
                f.write(f"{'='*60}\n")
                f.write(f"Generation {i}\n")
                f.write(f"{'='*60}\n")
                f.write(text + "\n\n")
        
        print(f"[Generator] Saved to {output_path}")

    def _save_tinygsm_results(self, results: Dict):
        """保存 TinyGSM/GSM8K 评估结果"""
        # 使用实例属性 self.output_dir
        output_dir = Path(self.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        full_path = output_dir / "gsm8k_results.txt"
        with open(full_path, 'w', encoding='utf-8') as f:
            for i, item in enumerate(results['full_outputs'], 1):
                f.write(f"{'='*60}\n")
                f.write(f"Problem {i}\n")
                f.write(f"{'='*60}\n")
                f.write(f"Question:\n{item['question']}\n\n")
                f.write(f"Solution:\n{item['solution']}\n\n")
                
                if item.get('code'):
                    f.write(f"Extracted Code:\n{item['code']}\n\n")
                
                f.write(f"Answer: {item['answer']}\n")
                
                if 'ground_truth' in results and i <= len(results['ground_truth']):
                    truth = results['ground_truth'][i-1]
                    # 由于我们在 _generate_solutions 中已经比较了，这里的比较用于最终文件的标记
                    match = "✓" if self._compare_answers(item['answer'], truth) else "✗"
                    f.write(f"Truth: {truth} {match}\n")
                
                f.write(f"{'='*60}\n\n")
        
        print(f"[Generator] Saved results to {full_path}")
        
        if 'verification' in results:
            verify_path = output_dir / "verification.txt"
            with open(verify_path, 'w', encoding='utf-8') as f:
                v = results['verification']
                f.write(f"{'='*60}\n")
                f.write("GSM8K Verification Results\n")
                f.write(f"{'='*60}\n")
                f.write(f"Total: {v['total']}\n")
                f.write(f"Correct: {v['correct']}\n")
                f.write(f"Incorrect: {v['incorrect']}\n")
                f.write(f"No Answer: {v['no_answer']}\n")
                f.write(f"Accuracy: {v['accuracy']*100:.2f}%\n")
                f.write(f"{'='*60}\n\n")
                
                for detail in v['details']:
                    status = "✓" if detail['correct'] else "✗"
                    f.write(f"{detail['index']+1}. {status} ")
                    f.write(f"Pred: {detail['predicted']}, Truth: {detail['ground_truth']}\n")
             
            print(f"[Generator] Saved verification to {verify_path}")