Create mat-gen.py

mat-gen.py

@@ -0,0 +1,373 @@
+import random
+import json
+import math
+import time
+import re
+import sys
+import multiprocessing
+import os
+from tqdm import tqdm
+
+NUM_LINES = 2000000
+OUTPUT_FILE = "correct_math_data.jsonl"
+
+MIN_LENGTH = 2
+MAX_LENGTH = 8
+MIN_NUMBER = 1
+MAX_NUMBER = 999
+MAX_EXPONENT_BASE = 9
+MAX_EXPONENT_POWER = 5
+
+REASONING_CHANCE = 0.8
+WORD_FORM_CHANCE = 0.25
+BRACKET_CHANCE = 0.5
+SENTENCE_FORM_CHANCE = 0.6
+MAX_SOLVER_ITERATIONS = 30  # Reduced from 50 for faster timeout
+
+NUM_WORKERS = os.cpu_count() or 1
+
+PROMPT_TEMPLATES = [
+    "What is {expression}?", "Calculate the value of {expression}.", "Find the result of {expression}.",
+    "Can you solve {expression}?", "Solve for {expression}.", "What does {expression} equal?", "Compute {expression}.",
+    "What is the solution to {expression}?", "Give me the answer for {expression}.", "Determine the value of {expression}.",
+    "Evaluate the expression: {expression}.", "I need the result of {expression}, please."
+]
+COT_INTRO_TEMPLATES = [
+    "<think> Let's break down the equation {expression} step by step, following the order of operations (BEDMAS).",
+    "<think> Okay, to solve {expression}, I'll follow BEDMAS (Brackets, Exponents, Division/Multiplication, Addition/Subtraction).",
+    "<think> Analyzing {expression}. I need to solve this by applying the correct order of operations.",
+    "<think> Here's my step-by-step evaluation for {expression}:",
+    "<think> To get the answer for {expression}, I will use the order of operations.",
+    "<think> Processing {expression} requires following BEDMAS, let's begin.",
+    "<think> I will solve {expression} by carefully following the rules of BEDMAS.",
+    "<think> The expression is {expression}. My plan is to solve it using the order of operations.",
+    "<think> To solve this, I'll go through Brackets, then Exponents, then Multiplication/Division, and finally Addition/Subtraction for {expression}.",
+    "<think> Let's start solving {expression}. I'll tackle it one operation at a time based on BEDMAS.",
+    "<think> Thinking step-by-step for {expression}..."
+]
+COT_STEP_TEMPLATES = {
+    "brackets": [
+        "First, I'll solve the expression inside the brackets: {part}. That equals {result}.",
+        "Starting with the parentheses, {part} evaluates to {result}.",
+        "The brackets are the priority. Calculating {part} gives me {result}.",
+        "The calculation inside the parentheses comes first: {part} becomes {result}.",
+        "Looking inside the brackets, I see {part}. The result of that is {result}.",
+        "I'll begin by simplifying the part in the parentheses: {part} is {result}.",
+        "The first step according to BEDMAS is brackets. So, {part} is solved to {result}.",
+        "Tackling the parentheses first: {part} simplifies to {result}.",
+        "Evaluating the bracketed expression {part} yields {result}.",
+        "My focus is on the brackets first. {part} equals {result}."
+    ],
+    "exponents": [
+        "Next, I'll handle the exponents. {part} is {result}.",
+        "Exponents are next in order. {part} calculates to {result}.",
+        "Now for the powers: {part} equals {result}.",
+        "Moving on to exponents, {part} results in {result}.",
+        "The next priority is exponents. The term {part} becomes {result}.",
+        "After brackets, I solve for exponents. {part} gives {result}.",
+        "Now, calculating the power: {part} is equal to {result}.",
+        "I see an exponent at {part}. This evaluates to {result}.",
+        "The 'E' in BEDMAS is for exponents, so I'll solve {part} to get {result}.",
+        "Time to resolve the exponents. {part} is {result}."
+    ],
+    "multi_div_mod": [
+        "Now, I'll perform multiplication, division, and modulo from left to right. The first is {part}, which is {result}.",
+        "Next up is multiplication and division. I see {part}, which gives {result}.",
+        "Working through multiplication/division from left to right, {part} results in {result}.",
+        "The next step is to resolve multiplication and division. {part} is {result}.",
+        "Scanning from left to right for M/D/M, I find {part}. This calculates to {result}.",
+        "Now for multiplication and division. The operation {part} equals {result}.",
+        "Moving on, I'll handle the multiplication/division. {part} becomes {result}.",
+        "The next operations are multiply and divide. I'll solve {part} to get {result}.",
+        "I will now compute {part}, which results in {result}.",
+        "Left-to-right, the next multiplication or division is {part}, giving {result}."
+    ],
+    "add_sub": [
+        "Finally, I'll do the addition and subtraction from left to right. I have {part}, which equals {result}.",
+        "Last step is addition and subtraction. {part} becomes {result}.",
+        "Finishing up with addition/subtraction, {part} evaluates to {result}.",
+        "The final operations are addition and subtraction. {part} results in {result}.",
+        "Now for the final calculations, addition and subtraction. {part} is {result}.",
+        "Working from left to right, the final step is {part}, which is {result}.",
+        "The last part of BEDMAS is addition and subtraction. {part} gives {result}.",
+        "To finish, I'll solve {part}, resulting in {result}.",
+        "Finally, the addition/subtraction part: {part} equals {result}.",
+        "The last calculation is {part}, and the answer is {result}."
+    ]
+}
+COT_FINALIZER_TEMPLATES = [
+    "After all steps, the final answer is {result}.",
+    "So, the complete result for the expression is {result}.",
+    "Therefore, the final value is {result}.",
+    "Bringing it all together, the answer is {result}.",
+    "The final computation yields {result}.",
+    "Thus, the expression evaluates to {result}.",
+    "So the final answer is {result}.",
+    "After all those steps, we arrive at the answer: {result}.",
+    "The result of the entire calculation is {result}.",
+    "In conclusion, the answer is {result}."
+]
+SIMPLE_COMPLETION_TEMPLATES = [
+    "The equation {expression} equals {result}.", "The answer is {result}.",
+    "The result is {result}.", "It equals {result}.", "The final value is {result}.",
+    "{expression} results in {result}.", "The solution is {result}.",
+    "The value is {result}.", "After calculation, the answer is {result}.",
+    "The final result is {result}."
+]
+
+ONES = ['', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine']
+TENS = ['', '', 'twenty', 'thirty', 'forty', 'fifty', 'sixty', 'seventy', 'eighty', 'ninety']
+TEENS = ['ten', 'eleven', 'twelve', 'thirteen', 'fourteen', 'fifteen', 'sixteen', 'seventeen', 'eighteen', 'nineteen']
+
+def number_to_words(n):
+    if not isinstance(n, int): return str(n)
+    if n == 0: return 'zero'
+    if n < 0: return f"negative {number_to_words(abs(n))}"
+    if n < 10: return ONES[n]
+    if n < 20: return TEENS[n-10]
+    if n < 100: return TENS[n//10] + (f"-{ONES[n%10]}" if n%10 else "")
+    if n < 1000: return f"{ONES[n//100]} hundred" + (f" and {number_to_words(n%100)}" if n%100 else "")
+    if n < 1000000: return f"{number_to_words(n//1000)} thousand" + (f", {number_to_words(n%1000)}" if n%1000 else "")
+    return str(n)
+
+def operator_to_word(op):
+    return {'+': 'plus', '-': 'minus', '*': 'times', '/': 'divided by', '^': 'to the power of', '%': 'modulo'}.get(op, op)
+
+def format_number(n):
+    if isinstance(n, float) and not n.is_integer():
+        return f"{n:.4f}".rstrip('0').rstrip('.')
+    return str(int(round(n)))
+
+def generate_expression_parts():
+    length = random.randint(MIN_LENGTH, MAX_LENGTH)
+    parts = []
+    for i in range(length):
+        if parts and parts[-1] == '^':
+            parts.append(random.randint(2, MAX_EXPONENT_POWER))
+        else:
+            parts.append(random.randint(MIN_NUMBER, MAX_NUMBER))
+
+        if i < length - 1:
+            if parts and parts[-1] != '^':
+                op = random.choice(['+', '-', '*', '/', '%', '^'])
+            else:
+                op = random.choice(['+', '-', '*', '/', '%'])
+
+            if op == '^':
+                parts[-1] = random.randint(MIN_NUMBER, MAX_EXPONENT_BASE)
+            parts.append(op)
+
+    if random.random() < BRACKET_CHANCE and len(parts) >= 5:
+        start = random.randrange(0, len(parts) - 2, 2)
+        end = random.randrange(start + 2, len(parts), 2)
+        parts.insert(end + 1, ')')
+        parts.insert(start, '(')
+    return parts
+
+def solve_with_cot(expression_str):
+    """Optimized solver with better pattern matching and guaranteed termination."""
+    steps = []
+    current_expr = expression_str.strip()
+    
+    for iteration in range(MAX_SOLVER_ITERATIONS):
+        # Remove extra spaces
+        current_expr = re.sub(r'\s+', ' ', current_expr).strip()
+        
+        # Check if we're done (single number)
+        try:
+            final_result = float(current_expr)
+            return {'steps': steps, 'result': final_result}
+        except ValueError:
+            pass
+        
+        reduction_made = False
+
+        # 1. Handle brackets first
+        bracket_match = re.search(r'\(([^()]+)\)', current_expr)
+        if bracket_match:
+            bracket_content = bracket_match.group(1).strip()
+            sub_solver_result = solve_with_cot(bracket_content)
+            if not sub_solver_result: 
+                return None
+            
+            result = sub_solver_result['result']
+            try:
+                formatted_result = format_number(result)
+            except (ValueError, OverflowError): 
+                return None
+
+            steps.append(random.choice(COT_STEP_TEMPLATES["brackets"]).format(part=bracket_content, result=formatted_result))
+            current_expr = current_expr[:bracket_match.start()] + ' ' + formatted_result + ' ' + current_expr[bracket_match.end():]
+            reduction_made = True
+            continue
+
+        # 2. Handle exponents
+        exp_match = re.search(r'(-?\d+(?:\.\d+)?)\s*\^\s*(-?\d+(?:\.\d+)?)', current_expr)
+        if exp_match:
+            base_str, exp_str = exp_match.groups()
+            try:
+                base = float(base_str)
+                exponent = float(exp_str)
+                result = base ** exponent
+                if abs(result) > 1e12 or math.isnan(result) or math.isinf(result):
+                    return None
+                formatted_result = format_number(result)
+            except (OverflowError, ValueError, ZeroDivisionError):
+                return None
+
+            part = f"{base_str} ^ {exp_str}"
+            steps.append(random.choice(COT_STEP_TEMPLATES["exponents"]).format(part=part, result=formatted_result))
+            current_expr = current_expr[:exp_match.start()] + ' ' + formatted_result + ' ' + current_expr[exp_match.end():]
+            reduction_made = True
+            continue
+
+        # 3. Handle multiplication, division, modulo (left to right)
+        mdm_match = re.search(r'(-?\d+(?:\.\d+)?)\s*([*/%])\s*(-?\d+(?:\.\d+)?)', current_expr)
+        if mdm_match:
+            left_str, op, right_str = mdm_match.groups()
+            try:
+                left = float(left_str)
+                right = float(right_str)
+                if op == '*':
+                    result = left * right
+                elif op == '/':
+                    if right == 0:
+                        return None
+                    result = left / right
+                elif op == '%':
+                    if right == 0:
+                        return None
+                    result = left % right
+                
+                if abs(result) > 1e12 or math.isnan(result) or math.isinf(result):
+                    return None
+                formatted_result = format_number(result)
+            except (OverflowError, ValueError, ZeroDivisionError):
+                return None
+
+            part = f"{left_str} {op} {right_str}"
+            steps.append(random.choice(COT_STEP_TEMPLATES["multi_div_mod"]).format(part=part, result=formatted_result))
+            current_expr = current_expr[:mdm_match.start()] + ' ' + formatted_result + ' ' + current_expr[mdm_match.end():]
+            reduction_made = True
+            continue
+
+        # 4. Handle addition and subtraction (left to right)
+        # Match pattern where we have number [+|-] number but not at start of negative number
+        as_match = re.search(r'(-?\d+(?:\.\d+)?)\s*([+\-])\s*(-?\d+(?:\.\d+)?)', current_expr)
+        if as_match:
+            left_str, op, right_str = as_match.groups()
+            try:
+                left = float(left_str)
+                right = float(right_str)
+                if op == '+':
+                    result = left + right
+                elif op == '-':
+                    result = left - right
+                
+                if abs(result) > 1e12 or math.isnan(result) or math.isinf(result):
+                    return None
+                formatted_result = format_number(result)
+            except (OverflowError, ValueError):
+                return None
+
+            part = f"{left_str} {op} {right_str}"
+            steps.append(random.choice(COT_STEP_TEMPLATES["add_sub"]).format(part=part, result=formatted_result))
+            current_expr = current_expr[:as_match.start()] + ' ' + formatted_result + ' ' + current_expr[as_match.end():]
+            reduction_made = True
+            continue
+        
+        # If no reduction was made, we're stuck - return None
+        if not reduction_made:
+            return None
+    
+    # Timeout reached
+    return None
+
+def generate_training_example(_=None):
+    """Generate a single training example with retry logic."""
+    max_retries = 50  # Reduced from 100 for faster generation
+    for attempt in range(max_retries):
+        try:
+            expression_parts = generate_expression_parts()
+            expression_str = " ".join(map(str, expression_parts))
+            
+            cot_result = solve_with_cot(expression_str)
+            
+            if cot_result and isinstance(cot_result['result'], (int, float)):
+                final_result = cot_result['result']
+                
+                # Filter out extreme values
+                if abs(final_result) > 1e12 or (final_result != 0 and abs(final_result) < 1e-4):
+                    continue
+                if math.isnan(final_result) or math.isinf(final_result):
+                    continue
+
+                result_str = format_number(final_result)
+                
+                if len(result_str) > 20: 
+                    continue
+
+                use_words = random.random() < WORD_FORM_CHANCE
+                if use_words:
+                    expression_text = ' '.join([number_to_words(p) if isinstance(p, int) else operator_to_word(p) if isinstance(p, str) else str(p) for p in expression_parts])
+                    result_text = number_to_words(int(round(final_result)))
+                    completion = random.choice(SIMPLE_COMPLETION_TEMPLATES).format(expression=expression_text, result=result_text)
+                else:
+                    expression_text = expression_str
+                    result_text = result_str
+                    use_reasoning = random.random() < REASONING_CHANCE
+                    if use_reasoning:
+                        intro = random.choice(COT_INTRO_TEMPLATES).format(expression=expression_text)
+                        steps_text = " ".join(cot_result['steps'])
+                        finalizer = random.choice(COT_FINALIZER_TEMPLATES).format(result=result_text)
+                        completion = f"{intro} {steps_text} {finalizer} </think>"
+                    else:
+                        completion = random.choice(SIMPLE_COMPLETION_TEMPLATES).format(expression=expression_text, result=result_text)
+
+                if random.random() < SENTENCE_FORM_CHANCE:
+                    prompt = random.choice(PROMPT_TEMPLATES).format(expression=expression_text)
+                else:
+                    prompt = f"{expression_text} ="
+
+                # Clean up spacing
+                prompt = re.sub(r'\s*\(', ' (', prompt)
+                prompt = re.sub(r'\)\s*', ') ', prompt).strip()
+                prompt = re.sub(r'\s+', ' ', prompt)
+                completion = re.sub(r'\s*\(', ' (', completion)
+                completion = re.sub(r'\)\s*', ') ', completion).strip()
+                completion = re.sub(r'\s+', ' ', completion)
+
+                return {"prompt": prompt, "completion": " " + completion}
+        except Exception as e:
+            continue
+            
+    return None
+
+def main():
+    print(f"🔥 Generating {NUM_LINES:,} examples using {NUM_WORKERS} parallel workers...")
+    print(f"   Appending to '{OUTPUT_FILE}'...")
+    start_time = time.time()
+    
+    generated_count = 0
+    failed_count = 0
+    
+    with open(OUTPUT_FILE, "a", encoding="utf-8") as f:
+        with multiprocessing.Pool(processes=NUM_WORKERS) as pool:
+            results_iterator = pool.imap_unordered(generate_training_example, range(NUM_LINES), chunksize=100)
+            
+            for item in tqdm(results_iterator, total=NUM_LINES, desc="Generating examples"):
+                if item:
+                    f.write(json.dumps(item) + "\n")
+                    generated_count += 1
+                else:
+                    failed_count += 1
+
+    elapsed_time = time.time() - start_time
+    print(f"\n\n✅ Done! Appended {generated_count:,} new items to '{OUTPUT_FILE}' in {elapsed_time:.2f}s.")
+    print(f"   📊 Success rate: {generated_count}/{NUM_LINES} ({100*generated_count/NUM_LINES:.1f}%)")
+    if failed_count > 0:
+        print(f"   ⚠️ {failed_count:,} generation attempts failed (expressions too complex or invalid)")
+
+if __name__ == "__main__":
+    multiprocessing.freeze_support()
+    main()