src/math_merge_eval.py

import os
# os.environ["VLLM_USE_V1"] = "0"
# os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
os.environ["VLLM_NO_USAGE_STATS"] = "1"
import json
import re
import sys
import gc
import random
import argparse
import traceback
from datetime import datetime
from typing import List, Dict, Optional, Any
from multiprocessing import Process, Queue

import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
from smpeft import PeftModel
from vllm import LLM, SamplingParams, EngineArgs
from datasets import load_dataset
import draccus
from tqdm import tqdm

try:
    from inference_math.grader import math_equal
except ImportError:
    raise ValueError("[Warning] 'grader.py' not found. GSM8k evaluation might fail.")

try:
    from inference_math import util
except ImportError:
    raise ValueError("[Warning] 'util.py' not found. MATH evaluation might fail.")

from .config import MainConfig 
from .utils import set_seed_all

MAX_NEW_TOKENS = 1536
MAX_INT = sys.maxsize

# CoT Prompt Template for Math
PROMPT_TEMPLATE = (
    "Below is an instruction that describes a task. "
    "Write a response that appropriately completes the request.\n\n"
    "### Instruction:\n{instruction}\n\n### Response: Let's think step by step."
)

# --- Helper Functions for GSM8k ---
from fraction import Fraction

def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        pass
    try:
        import unicodedata
        unicodedata.numeric(s)
        return True
    except (TypeError, ValueError):
        pass
    return False

def extract_answer_number_gsm8k(completion):
    """
    Exact logic from gsm8k_infer.py
    """
    text = completion.split('The answer is: ')
    if len(text) > 1:
        extract_ans = text[-1].strip()
        match = re.search(r'[\-+]?\d*[\.,/]?\d+', extract_ans)
        if match:
            if '/' in match.group():
                denominator = match.group().split('/')[1]
                numerator = match.group().split('/')[0]
                if is_number(denominator) == True and is_number(numerator) == True:
                    if denominator == '0':
                        return round(float(numerator.replace(',', '')))
                    else:
                        frac = Fraction(match.group().replace(',', ''))
                        num_numerator = frac.numerator
                        num_denominator = frac.denominator
                        return round(float(num_numerator / num_denominator))
                else:
                    return None
            else:
                if float(match.group().replace(',', '')) == float('inf'):
                    return None
                return round(float(match.group().replace(',', '')))
        else:
            return None
    else:
        return None
    
# --- Helper Functions: MATH (Ported strictly from MATH_infer.py) ---
def remove_boxed(s):
    """
    Extracts content from \boxed{...}
    """
    left = "\\boxed{"
    try:
        assert s[:len(left)] == left
        assert s[-1] == "}"
        return s[len(left):-1]
    except:
        return None
    
def process_results_math(completion, answer):
    """
    Exact logic from MATH_infer.py
    """
    split_ans = completion.split('The answer is: ')
    if len(split_ans) > 1:
        ans = split_ans[-1]
        extract_ans_temp = ans.split('.\n')[0]
        extract_ans_temp = extract_ans_temp.strip()
        if len(extract_ans_temp)>0 and extract_ans_temp[-1] == '.':
            extract_ans = extract_ans_temp[0:-1]
        else:
            extract_ans = extract_ans_temp
        extract_ans = extract_ans.strip()
        
        if util.is_equiv(extract_ans, answer):
            return True, extract_ans
        else:
            return False, extract_ans
    else:
        return False, None
    
# --- Prompt Formatting ---
def format_prompt(examples):
    prompts = []
    # GSM8k uses 'question', MATH often uses 'instruction' or 'problem'
    instructions = examples.get('question', examples.get('instruction', []))
    
    for instr in instructions:
        source_text = PROMPT_TEMPLATE.format(instruction=instr)
        prompts.append(source_text)
    
    return {"prompt": prompts}

# --- Core Logic ---
def merge_process(queue, mainCfg: MainConfig, force_to_merge: bool = False):
    """
    Handles the PEFT merge process in a separate process to manage VRAM.
    """
    try:
        model_name = mainCfg.model.model_name
        
        # Determine adapter path
        if mainCfg.model.merge_adapter_path is not None:
            adapter = mainCfg.model.merge_adapter_path + "/ft2" # Adjust subfolder if needed
            print(f'Merging from merge_adapter_path: {adapter}')
        elif mainCfg.model.adapter_path is not None:
            adapter = mainCfg.model.adapter_path + "/ft2"
            print(f'Merging from adapter_path: {adapter}')
        else:
            raise KeyError('No adapter path provided in config.')
        
        # Determine output path
        if mainCfg.model.merge_output_path is not None:
            output_path = os.path.join(mainCfg.model.merge_output_path, "merge")
            out_json = mainCfg.model.merge_output_path
        else:
            output_path = os.path.join(mainCfg.model.adapter_path, "merge")
            out_json = mainCfg.model.adapter_path
        
        # Check if merge is needed
        if os.path.exists(output_path):
            has_weights = any(f.endswith(".bin") or f.endswith(".safetensors") for f in os.listdir(output_path))
        else:
            has_weights = False

        if not has_weights or force_to_merge:
            print(f"Loading base model: {model_name}")
            model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", low_cpu_mem_usage=True)
            tokenizer = AutoTokenizer.from_pretrained(model_name, device_map='cpu')

            print(f"Loading adapter: {adapter}")
            model = PeftModel.from_pretrained(model, adapter)
            print("Merging model...")
            model = model.merge_and_unload()
            
            print(f"Saving merged model to: {output_path}")
            model.save_pretrained(output_path, safe_serialization=True, max_shard_size="10GB")
            tokenizer.save_pretrained(output_path)
            
            del model
            del tokenizer
            gc.collect()
            torch.cuda.empty_cache()
            print('Merge complete.')
        else:
            print("Merged weights found. Skipping merge step.")
            
        queue.put((output_path, out_json))
    
    except Exception as e:
        error_msg = traceback.format_exc()
        print(error_msg)
        queue.put(error_msg)
        print(f"Error in merge_process: {e}")

# --- Scoring Logic ---
def score_outputs(outputs, test_target_name, ground_truths, out_json):
    results = []
    total_correct = 0
    total_samples = len(ground_truths)
    invalid_count = 0

    print(f"Calculating scores for {test_target_name}...")

    # Identify dataset type
    is_gsm8k = 'gsm8k' in test_target_name.lower()
    is_math = 'math' in test_target_name.lower()

    for i, output in enumerate(tqdm(outputs, desc="Scoring")):
        prediction_text = output.outputs[0].text
        gt_raw = ground_truths[i]
        
        is_correct = False
        extracted_pred = None
        clean_gt = None

        # --- GSM8k Scoring ---
        if is_gsm8k:
            # Logic: GT format in jsonl is usually "Reasoning.... #### 1234"
            # gsm8k_infer.py logic: temp_ans = item['answer'].split('#### ')[1]; int(replace(','))
            try:
                if '####' in str(gt_raw):
                    clean_gt_str = str(gt_raw).split('#### ')[1].replace(',', '').strip()
                else:
                    clean_gt_str = str(gt_raw).replace(',', '').strip()
                
                clean_gt = float(clean_gt_str)
            except:
                clean_gt = gt_raw # Keep raw if parsing fails

            extracted_pred = extract_answer_number_gsm8k(prediction_text)
            
            if extracted_pred is not None:
                # Comparison logic from gsm8k_infer.py
                try:
                    is_correct = (float(extracted_pred) == float(clean_gt)) or math_equal(extracted_pred, clean_gt)
                except:
                    is_correct = False
            else:
                is_correct = False
                invalid_count += 1

        # --- MATH Scoring ---
        elif is_math:
            # Logic: GT in jsonl is usually the solution text
            # MATH_infer.py logic: remove_boxed(util.last_boxed_only_string(solution))
            try:
                clean_gt = remove_boxed(util.last_boxed_only_string(str(gt_raw)))
            except:
                clean_gt = gt_raw

            # Logic from MATH_infer.py: process_results
            is_correct, extracted_pred = process_results_math(prediction_text, clean_gt)
            
            if not extracted_pred and not is_correct:
                invalid_count += 1

        results.append({
            "id": i,
            "prediction_full": prediction_text,
            "extracted_pred": extracted_pred,
            "ground_truth_raw": gt_raw,
            "ground_truth_clean": clean_gt,
            "is_correct": is_correct,
        })
        
        if is_correct:
            total_correct += 1

    avg_acc = 100.0 * total_correct / total_samples if total_samples > 0 else 0
    
    print("\n" + "="*40)
    print(f"FINAL RESULTS: {test_target_name}")
    print("="*40)
    print(f"Total Samples: {total_samples}")
    print(f"Invalid/No Answer Found: {invalid_count}")
    print(f"Accuracy: {avg_acc:.2f}%")
    print("="*40)

    os.makedirs(out_json, exist_ok=True)
    save_file = os.path.join(out_json, f'{test_target_name}.json')
    with open(save_file, "w", encoding="utf-8") as f:
        json.dump({
            "metrics": {
                "accuracy": avg_acc,
                "total": total_samples,
                "invalid": invalid_count
            },
            "details": results
        }, f, indent=2, ensure_ascii=False)
        
    return avg_acc

@draccus.wrap()
def main(mainCfg: MainConfig):
    print('='*120)
    set_seed_all(mainCfg.seed)
    
    # --- 1. Merge Step ---
    queue = Queue()
    p = Process(target=merge_process, args=(queue, mainCfg, False))
    p.start()
    merge_result = queue.get() 
    p.join()

    if merge_result is None or isinstance(merge_result, str): # Handle error string
        raise RuntimeError(f"Model merging failed: {merge_result}")
    
    model_path, out_json = merge_result

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model directory does not exist: {model_path}")
    
    print(f"Verified model path: {os.path.abspath(model_path)}")
    out_json = os.path.join(out_json, "results")
    print('Output JSON path: ', out_json)

    # --- 2. Initialize vLLM ---
    print("Initializing vLLM...")
    llm = LLM(
        model=model_path,
        dtype="bfloat16",
        gpu_memory_utilization=0.9, 
        max_model_len=mainCfg.infer.infer_max_seq_length,
        tensor_parallel_size=1,
    )

    # Stop tokens for CoT
    stop_tokens = ["Instruction:", "Instruction", "Response:", "Response",]
    sampling_params = SamplingParams(
        temperature=0, 
        top_p=1, 
        max_tokens=MAX_NEW_TOKENS, 
        stop=stop_tokens
    )
    
    start_time_total = datetime.now()
    final_res = {}
    all_task_acc = []

    # --- 3. Evaluation Loop ---
    try:
        # Loop through datasets defined in config
        for test_target_name in mainCfg.infer.datasets:
            print(f"Processing dataset: {test_target_name}")
            
            # Load Local Dataset
            dataset_path = f'./dataset/{test_target_name}/test.jsonl'
            if not os.path.exists(dataset_path):
                print(f"[Error] Local file not found: {dataset_path}. Skipping.")
                continue

            print(f"Loading local file: {dataset_path}")
            test_dataset = load_dataset("json", data_files=dataset_path, split='train')

            # Standardize Column Names
            # MATH: instruction -> question, output -> answer
            if 'instruction' in test_dataset.column_names:
                test_dataset = test_dataset.rename_column('instruction', 'question')
            if 'output' in test_dataset.column_names:
                test_dataset = test_dataset.rename_column('output', 'answer')

            ground_truths = test_dataset['answer']

            # Format Prompts
            print("Formatting prompts...")
            test_dataset = test_dataset.map(
                format_prompt,
                batched=True,
                batch_size=1000,
                desc="Formatting prompts"
            )
            prompts = test_dataset['prompt']

            # Generate
            print(f"Generating responses for {len(prompts)} samples...")
            start_time_task = datetime.now()
            
            outputs = llm.generate(prompts, sampling_params)
            
            end_time_task = datetime.now()
            print(f"Task {test_target_name} duration: {end_time_task - start_time_task}")
            
            # Score
            avg_acc = score_outputs(
                outputs=outputs,
                test_target_name=test_target_name,
                ground_truths=ground_truths,
                out_json=out_json
            )
            
            final_res[test_target_name] = avg_acc
            all_task_acc.append(avg_acc)
            
            del prompts
            del outputs
            del test_dataset
            # gc.collect()

    except Exception as e:
        print(f"Error during evaluation loop: {e}")
        traceback.print_exc()

    # --- 4. Final Report ---
    print('Accuracies per task:', all_task_acc)
    if all_task_acc:
        avg_score = sum(all_task_acc) / len(all_task_acc)
    else:
        avg_score = 0.0

    final_res['average_score'] = avg_score
    
    os.makedirs(out_json, exist_ok=True)
    save_file = os.path.join(out_json, 'FINAL.json')
    
    with open(save_file, "w", encoding="utf-8") as f:
        json.dump(final_res, f, indent=2, ensure_ascii=False)
        
    print(f"All Results saved to {save_file}, Overall Score: {avg_score:.2f}")

    end_time_total = datetime.now()
    print(f"Total execution time: {end_time_total - start_time_total}")

if __name__ == "__main__":
    # ------------------------------------------------------------------
    # EXPLANATION:
    # We must force the start method to 'spawn'.
    # Why? By default, Linux uses 'fork'. If the parent process has 
    # already initialized CUDA (even by just checking availability or 
    # running inside a VS Code debugger), 'fork' will copy that corrupted 
    # CUDA context to the child process, causing a RuntimeError.
    # 'spawn' starts a fresh interpreter, avoiding this issue entirely.
    # ------------------------------------------------------------------
    # try:
    #     import torch.multiprocessing as mp
    #     mp.set_start_method('spawn', force=True)
    # except RuntimeError as e:
    #     # If it's already set, just print a warning (safe to ignore usually)
    #     print(f"Warning: context already set: {e}")
    main()