verifier.py

"""
Script to verify generated math-eval dataset outputs.

Usage:
    python verifier.py --file <equations_file>
"""
import re
import argparse
import os
import sys

def parse_equations(line):
    """Extract multiple equations and their respective solutions."""
    equation_part, solution_part = line.split("<sep>")

    # Extract variable assignments
    solution_dict = {}
    for assign in solution_part.strip().split(","):
        var, value = assign.split("=")
        solution_dict[var.strip()] = int(value.strip())

    # Split multiple equations by commas
    equations = equation_part.split(",")

    parsed_equations = []
    for eq in equations:
        eq = eq.strip()
        equation_match = re.findall(r'([-+]?\d*)\s*([a-zA-Z])', eq)
        constant_match = re.search(r'=\s*(-?\d+)', eq)

        coefficients = {}
        for coef, var in equation_match:
            coefficients[var] = int(coef) if coef not in ["", "+", "-"] else (1 if coef in ["", "+"] else -1)

        constant = int(constant_match.group(1)) if constant_match else 0
        parsed_equations.append((coefficients, constant))

    return parsed_equations, solution_dict

def verify_equations(parsed_equations, solution_dict):
    """Verify if all equations in a line satisfy the given solutions."""
    return all(
        sum(coefficients[var] * solution_dict[var] for var in coefficients) == constant
        for coefficients, constant in parsed_equations
    )

def verify_file(filename):
    """Read a text file and verify each equation set."""
    total_lines = 0
    successful_lines = 0
    
    with open(filename, "r") as file:
        for line_number, line in enumerate(file, start=1):
            line = line.strip()
            if not line:
                continue
                
            total_lines += 1
            try:
                parsed_equations, solution_dict = parse_equations(line)
                if verify_equations(parsed_equations, solution_dict):
                    successful_lines += 1
                else:
                    print(f"Line {line_number}: Verification failed for {line}")
            except Exception as e:
                print(f"Line {line_number}: Error parsing {line} - {e}")
    
    print(f"\nVerification Summary:")
    print(f"Total equations: {total_lines}")
    print(f"Successfully verified: {successful_lines}")
    print(f"Failed: {total_lines - successful_lines}")
    print(f"Success rate: {(successful_lines/total_lines*100):.1f}%" if total_lines > 0 else "No equations found")
    
    return successful_lines == total_lines

def main():
    parser = argparse.ArgumentParser(
        description="Verify generated math equations.",
        usage="python verifier.py --file <equations_file>"
    )
    parser.add_argument('--file', type=str, required=True, help='Path to the equations file to verify')
    args = parser.parse_args()
    
    if not os.path.exists(args.file):
        print(f"Error: File {args.file} not found")
        return False
    
    print(f"Verifying equations in: {args.file}")
    return verify_file(args.file)

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)