GSM8K.py

import os
import re
import sys
import pandas as pd
from typing import List, Dict, Tuple, Optional

current_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(current_dir)
sys.path.append(root_dir)


def load_gsm8k_data(file_path: str) -> List[Dict]:

    print(f"Loading GSM8K data for Llama3 from: {file_path}")
    try:
        df = pd.read_parquet(file_path)
        required_cols = ['question', 'answer']
        for col in required_cols:
            if col not in df.columns:
                raise ValueError(f"Parquet lost: {col}")
   
        records = df.to_dict('records')
        print(f"Successfully loaded {len(records)} samples.")
        return records
    except Exception as e:
        print(f"An error occurred while loading the data: {e}")
        sys.exit(1)

def extract_ground_truth(answer_str: str) -> Optional[float]:

    match = re.search(r"####\s*([-]?\d{1,3}(?:,\d{3})*\.?\d*)", answer_str)
    if match:
        gt_str = match.group(1).replace(",", "")
        try:
            return float(gt_str)
        except ValueError:
            return None
    return None


def build_gsm8k_prompt(record: dict, tokenizer) -> Tuple[str, Optional[float]]:

    question_text = record['question']
    ground_truth_text = record['answer']

    ground_truth_answer = extract_ground_truth(ground_truth_text)

    output_format_instruction = (
        'Your response *must* end with "The final answer is (answer)". '
        'No units. For example:\n(Question and your reasoning)\nThe final answer is 33.'
    )

    user_instruction = (
        f"Please answer the following question:\n\n"
        f"{question_text}\n\n"
        f"{output_format_instruction}"
    )

    messages = [
        {"role": "system",
         "content": "You are a helpful expert problem solver. Please strictly follow the user's instructions, especially the output format."},
        {"role": "user", "content": user_instruction}
    ]

    full_prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    return full_prompt, ground_truth_answer

def extract_gsm8k_answer(generation_only_text: str) -> Optional[float]:

    text = generation_only_text.strip()
    match = re.search(r"The final answer is\s*([-]?\d{1,3}(?:,\d{3})*\.?\d*)", text)
    if match:
        answer_str = match.group(1).replace(",", "")
        try:
            return float(answer_str)
        except ValueError:
            return None
    return None