mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-28 17:29:30 +00:00
word problem generation working
This commit is contained in:
Shannon Sands
parent
5c0c7f5b10
commit
fd5b87011d
2 changed files with 180 additions and 38 deletions
|
|
@ -1,15 +1,21 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from openai import AsyncOpenAI, NotGiven
|
||||
|
||||
from atroposlib.envs.base import BaseEnv, BaseEnvConfig, OpenaiConfig, ScoredDataGroup
|
||||
from atroposlib.utils.tokenize_for_trainer import tokenize_for_trainer
|
||||
|
||||
from .curriculum import MathCurriculum
|
||||
|
||||
load_dotenv()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
|
|
@ -34,7 +40,7 @@ Your answer format should be:
|
|||
|
||||
\\boxed{your final answer here}
|
||||
|
||||
Remember to format your final answer correctly as this is important for evaluation."""
|
||||
Remember to format your final answer correctly as this is important for evaluation. Do not apply any rounding to your final answer, be as exact as possible."""
|
||||
|
||||
|
||||
class InfiniteMathEnvConfig(BaseEnvConfig):
|
||||
|
|
@ -56,6 +62,11 @@ class InfiniteMathEnvConfig(BaseEnvConfig):
|
|||
temperature: float = 0.7
|
||||
top_p: float = 0.9
|
||||
|
||||
# Model for word problem generation
|
||||
word_problem_model_name: Optional[str] = "gpt-4.1-mini"
|
||||
word_problem_openai_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
|
||||
word_problem_openai_base_url: Optional[str] = None # Add for custom server
|
||||
|
||||
|
||||
class InfiniteMathEnv(BaseEnv):
|
||||
"""Environment for procedurally generated math problems with curriculum advancement."""
|
||||
|
|
@ -245,16 +256,116 @@ class InfiniteMathEnv(BaseEnv):
|
|||
|
||||
await super().wandb_log(wandb_metrics)
|
||||
|
||||
async def _convert_to_word_problem(self, raw_problem_text: str) -> str:
|
||||
"""Converts a raw math problem string into a word problem using an LLM."""
|
||||
system_prompt_word_problem = """You are an expert creative writer. Your task is to transform a given raw mathematical expression into an engaging and imaginative word problem.
|
||||
|
||||
**Critical Instructions:**
|
||||
1. **Strict Preservation:** The core mathematical question, ALL numbers, and ALL operations from the raw problem MUST be EXACTLY preserved in the word problem. Do NOT change the calculation required. For example, if the raw problem is 'A - B', the word problem must represent subtraction of B from A, not any other operation.
|
||||
2. **Clarity:** The word problem must clearly and unambiguously lead to solving the original mathematical expression.
|
||||
3. **Conciseness:** Keep the word problem relatively short and to the point.
|
||||
4. **Output Format:** Output ONLY the word problem text. Do NOT include any preambles, self-references (like 'Here is a word problem:'), special tokens (like '<|start_header_id|>'), or any text other than the word problem itself.
|
||||
|
||||
**Examples of Correct Transformation:**
|
||||
Raw Problem: 5 * 3
|
||||
Word Problem: Sarah is baking cookies, and each batch requires 3 eggs. If Sarah wants to bake 5 batches, how many eggs will she need in total?
|
||||
|
||||
Raw Problem: |10 - 15|
|
||||
Word Problem: A submarine is 10 meters below sea level. Another submarine is 15 meters below sea level. What is the absolute difference in their depths in meters?
|
||||
|
||||
Raw Problem: sqrt(16)
|
||||
Word Problem: A square piece of land has an area of 16 square units. What is the length of one of its sides in units?
|
||||
|
||||
**Example of Incorrect Transformation (Operation Changed):**
|
||||
Raw Problem: |3 - (-67)| (This is 3 + 67)
|
||||
Incorrect Word Problem: In a magical forest, there are 3 enchanted trees, and each tree has 67 glowing fruits. How many glowing fruits are there in total? (This became 3 * 67)
|
||||
Correct Word Problem: A bird watcher is 3 meters up a tree. She spots a rare bird 67 meters below ground level in a cave. What is the total vertical distance between the bird watcher and the rare bird in meters?
|
||||
"""
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt_word_problem},
|
||||
{"role": "user", "content": f"Raw Problem: {raw_problem_text}"},
|
||||
]
|
||||
|
||||
prompt_for_llm = self.tokenizer.apply_chat_template(messages, tokenize=False)
|
||||
|
||||
try:
|
||||
api_key_to_use = self.config.word_problem_openai_api_key or os.environ.get(
|
||||
"OPENAI_API_KEY"
|
||||
)
|
||||
base_url_to_use = self.config.word_problem_openai_base_url
|
||||
model_to_use = (
|
||||
self.config.word_problem_model_name or "gpt-4.1-mini"
|
||||
) # Fallback if not set in config somehow
|
||||
|
||||
if not api_key_to_use:
|
||||
logger.error(
|
||||
"OpenAI API key for word problem generation is not configured (checked config and OPENAI_API_KEY env var)."
|
||||
)
|
||||
return raw_problem_text # Fallback if no API key
|
||||
|
||||
client = AsyncOpenAI(
|
||||
api_key=api_key_to_use,
|
||||
base_url=base_url_to_use, # If base_url_to_use is None, client uses default. No need for NotGiven here.
|
||||
)
|
||||
|
||||
chat_completions = await client.chat.completions.create(
|
||||
model=model_to_use,
|
||||
messages=messages,
|
||||
n=1,
|
||||
max_tokens=512,
|
||||
temperature=0.7,
|
||||
top_p=1.0,
|
||||
)
|
||||
|
||||
generated_text = chat_completions.choices[0].message.content
|
||||
|
||||
original_llm_output = (
|
||||
generated_text # Store raw LLM output for logging if cleaning fails
|
||||
)
|
||||
|
||||
# Simplified Cleaning: Only strip whitespace now
|
||||
cleaned_text = original_llm_output.strip()
|
||||
|
||||
if (
|
||||
not cleaned_text
|
||||
): # If cleaning (now just stripping) results in an empty string, fallback
|
||||
logger.warning(
|
||||
f"Word problem conversion for '{raw_problem_text}' resulted in empty string after stripping. Original LLM output was: '{original_llm_output}'. Falling back to raw problem."
|
||||
)
|
||||
return raw_problem_text
|
||||
|
||||
logger.info(
|
||||
f"Converted raw problem '{raw_problem_text}' to word problem: '{cleaned_text}'"
|
||||
)
|
||||
return cleaned_text
|
||||
except Exception as e:
|
||||
log_message_error = (
|
||||
f"Error converting to word problem for '{raw_problem_text}': {e}"
|
||||
)
|
||||
logger.error(log_message_error)
|
||||
return raw_problem_text
|
||||
|
||||
async def get_next_item(self):
|
||||
"""Get the next problem based on current curriculum level."""
|
||||
problem, solution, generator_id = self.curriculum.get_problem()
|
||||
raw_problem, solution, generator_id = self.curriculum.get_problem()
|
||||
|
||||
problem = self._strip_latex_delimiters(problem)
|
||||
solution = self._strip_latex_delimiters(solution)
|
||||
# Strip LaTeX delimiters from the raw problem before converting to word problem
|
||||
raw_problem_stripped = self._strip_latex_delimiters(raw_problem)
|
||||
# Also strip from solution for consistency, though solution isn't used in word problem conversion
|
||||
solution_stripped = self._strip_latex_delimiters(solution)
|
||||
|
||||
prompt = tuple([frozenset({"role": "user", "content": problem}.items())])
|
||||
# Convert the stripped raw problem to a word problem
|
||||
word_problem_text = await self._convert_to_word_problem(raw_problem_stripped)
|
||||
|
||||
return (prompt, solution, generator_id)
|
||||
# Create a message with the word problem
|
||||
# The agent will solve this word problem, which should map back to the original solution.
|
||||
prompt = tuple(
|
||||
[frozenset({"role": "user", "content": word_problem_text}.items())]
|
||||
)
|
||||
|
||||
# Return the word problem with the original (stripped) solution and generator_id
|
||||
return (prompt, solution_stripped, generator_id)
|
||||
|
||||
async def evaluate(self, *args, **kwargs):
|
||||
"""Evaluate the model on test problems at the current curriculum level."""
|
||||
|
|
@ -316,14 +427,21 @@ class InfiniteMathEnv(BaseEnv):
|
|||
) -> Tuple[int, bool]:
|
||||
"""Evaluate a single problem."""
|
||||
try:
|
||||
logger.debug(f"Evaluating level {level} problem: {problem[:30]}...")
|
||||
# Problem here is already stripped of LaTeX by the setup method
|
||||
# Convert the raw problem to a word problem for evaluation
|
||||
word_problem_text = await self._convert_to_word_problem(problem)
|
||||
logger.debug(
|
||||
f"Evaluating level {level} word problem: {word_problem_text[:50]}... (Original raw: {problem[:30]}...)"
|
||||
)
|
||||
|
||||
# Convert messages to a single prompt using the tokenizer
|
||||
messages = [
|
||||
{"role": "system", "content": self.system_prompt},
|
||||
{"role": "user", "content": problem},
|
||||
{"role": "user", "content": word_problem_text}, # Use word problem here
|
||||
]
|
||||
prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
|
||||
|
||||
# Add prefilled thinking starter
|
||||
prefill = "\n<think>\n"
|
||||
prefilled_prompt = prompt + prefill
|
||||
|
||||
|
|
@ -514,15 +632,19 @@ class InfiniteMathEnv(BaseEnv):
|
|||
starting_level=1,
|
||||
progress_threshold=0.8,
|
||||
min_evaluations=10,
|
||||
max_attempts_per_problem=3, # Default from class, not in old main
|
||||
correct_reward=1.0, # As in old main
|
||||
incorrect_reward=-0.5, # As in old main (class default was -1.0)
|
||||
think_block_bonus=0.2, # As per previous update
|
||||
boxed_answer_bonus=0.2, # As per previous update
|
||||
apply_length_penalty=True, # As in old main
|
||||
length_threshold_ratio=0.6, # As in old main (class default was 0.5)
|
||||
temperature=0.7, # As in old main
|
||||
top_p=0.9, # As in old main
|
||||
max_attempts_per_problem=3,
|
||||
correct_reward=1.0,
|
||||
incorrect_reward=-0.5,
|
||||
think_block_bonus=0.2,
|
||||
boxed_answer_bonus=0.2,
|
||||
apply_length_penalty=True,
|
||||
length_threshold_ratio=0.6,
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
# Specify the model and connection details for word problem generation
|
||||
word_problem_model_name="gpt-4.1-mini",
|
||||
word_problem_openai_api_key=None, # Default to None (uses OPENAI_API_KEY env var)
|
||||
word_problem_openai_base_url=None, # Default to None (uses official OpenAI endpoint)
|
||||
)
|
||||
|
||||
server_configs = [
|
||||
|
|
|
|||
|
|
@ -43,6 +43,10 @@ async def main():
|
|||
length_threshold_ratio=0.6,
|
||||
temperature=0.3,
|
||||
top_p=0.9,
|
||||
word_problem_model_name="gpt-4.1-mini",
|
||||
word_problem_openai_api_key=os.getenv("OPENAI_API_KEY_WORD_PROBLEM")
|
||||
or os.getenv("OPENAI_API_KEY"),
|
||||
word_problem_openai_base_url=os.getenv("OPENAI_BASE_URL_WORD_PROBLEM"),
|
||||
)
|
||||
|
||||
server_configs = [
|
||||
|
|
@ -53,12 +57,11 @@ async def main():
|
|||
timeout=600,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
logger.info("Using hardcoded debug configuration.")
|
||||
logger.debug(f"Env Config: {config}")
|
||||
logger.debug(f"Server Configs: {server_configs}")
|
||||
|
||||
|
||||
try:
|
||||
env = InfiniteMathEnv(
|
||||
config=config,
|
||||
|
|
@ -77,19 +80,23 @@ async def main():
|
|||
item = await env.get_next_item()
|
||||
problem_prompt, solution, generator_id = item
|
||||
|
||||
problem_content = dict(problem_prompt[0])['content']
|
||||
logger.info(f"Problem (ID: {generator_id}, Level: {env.curriculum.get_current_level()}): {problem_content}")
|
||||
problem_content = dict(problem_prompt[0])["content"]
|
||||
logger.info(
|
||||
f"Problem (ID: {generator_id}, Level: {env.curriculum.get_current_level()}): {problem_content}"
|
||||
)
|
||||
logger.info(f"Expected Solution: {solution}")
|
||||
|
||||
logger.info("Collecting trajectories...")
|
||||
trajectories_data, backlog = await env.collect_trajectories(item)
|
||||
|
||||
|
||||
if not trajectories_data:
|
||||
logger.error("No trajectories were collected.")
|
||||
return
|
||||
|
||||
logger.info(f"Collected {len(trajectories_data)} data points for scoring (should be 1 for group_size=1).")
|
||||
|
||||
logger.info(
|
||||
f"Collected {len(trajectories_data)} data points for scoring (should be 1 for group_size=1)."
|
||||
)
|
||||
|
||||
logger.info("Scoring trajectories...")
|
||||
scored_data = await env.score(trajectories_data)
|
||||
|
||||
|
|
@ -99,7 +106,7 @@ async def main():
|
|||
assistant_response = ""
|
||||
if messages_list and messages_list[-1].get("role") == "assistant":
|
||||
assistant_response = messages_list[-1].get("content", "N/A")
|
||||
|
||||
|
||||
logger.info(f"--- Attempt {i+1} ---")
|
||||
logger.info(f"Problem: {problem_content}")
|
||||
logger.info(f"Full Assistant Response:\\n{assistant_response}")
|
||||
|
|
@ -107,21 +114,22 @@ async def main():
|
|||
is_correct_task = env.check_answer(assistant_response, solution)
|
||||
logger.info(f"Checked Correct by env.check_answer: {is_correct_task}")
|
||||
|
||||
|
||||
correct_count_buffer = sum(env.percent_correct_buffer)
|
||||
total_attempts_buffer = len(env.percent_correct_buffer)
|
||||
|
||||
|
||||
logger.info("\n--- Overall for this run ---")
|
||||
logger.info(f"Expected Solution: {solution}")
|
||||
logger.info(f"Score(s) from env.score: {scored_data['scores']}")
|
||||
if total_attempts_buffer > 0:
|
||||
logger.info(f"Correct based on internal buffer: {correct_count_buffer}/{total_attempts_buffer}")
|
||||
logger.info(
|
||||
f"Correct based on internal buffer: {correct_count_buffer}/{total_attempts_buffer}"
|
||||
)
|
||||
else:
|
||||
logger.info("No attempts recorded in percent_correct_buffer.")
|
||||
|
||||
else:
|
||||
logger.error("Scored data is missing expected fields ('messages' or 'scores').")
|
||||
|
||||
|
||||
logger.info("=======================================")
|
||||
|
||||
# Re-add curriculum and evaluation testing
|
||||
|
|
@ -148,10 +156,14 @@ async def main():
|
|||
# Check if the level advanced
|
||||
new_level_eval = env.curriculum.get_current_level()
|
||||
if new_level_eval > initial_level_eval:
|
||||
logger.info(f"Successfully advanced from level {initial_level_eval} to level {new_level_eval} during evaluation!")
|
||||
logger.info(
|
||||
f"Successfully advanced from level {initial_level_eval} to level {new_level_eval} during evaluation!"
|
||||
)
|
||||
logger.info(f"New level description: {env.curriculum.get_level_description()}")
|
||||
else:
|
||||
logger.info(f"Did not advance during evaluation. Remained at level {initial_level_eval}.")
|
||||
logger.info(
|
||||
f"Did not advance during evaluation. Remained at level {initial_level_eval}."
|
||||
)
|
||||
# Show current progress toward advancement
|
||||
current_level_desc = env.curriculum.get_current_level()
|
||||
if current_level_desc in env.curriculum.performance_history:
|
||||
|
|
@ -182,20 +194,24 @@ async def main():
|
|||
# Test curriculum advancement with simulated performance history
|
||||
logger.info("\n=== Testing Curriculum Advancement Manually ===")
|
||||
initial_level_manual_adv = env.curriculum.get_current_level()
|
||||
logger.info(f"Starting manual advancement test from level: {initial_level_manual_adv}")
|
||||
logger.info(
|
||||
f"Starting manual advancement test from level: {initial_level_manual_adv}"
|
||||
)
|
||||
|
||||
# Simulate good performance at current level
|
||||
# Ensure we don't try to get items if curriculum is already at max level from previous eval
|
||||
max_level_possible = max(env.curriculum.DIFFICULTY_LEVELS.keys())
|
||||
if initial_level_manual_adv < max_level_possible:
|
||||
logger.info(f"Simulating {config.min_evaluations} correct answers for level {initial_level_manual_adv}...")
|
||||
for _ in range(config.min_evaluations): # Use config for min_evaluations
|
||||
logger.info(
|
||||
f"Simulating {config.min_evaluations} correct answers for level {initial_level_manual_adv}..."
|
||||
)
|
||||
for _ in range(config.min_evaluations): # Use config for min_evaluations
|
||||
# Get a problem from current level to ensure generator_id is valid for the level
|
||||
# The level might have changed due to the previous env.evaluate() call
|
||||
problem_item_adv_test = await env.get_next_item()
|
||||
problem_item_adv_test = await env.get_next_item()
|
||||
_, _, generator_id_adv_test = problem_item_adv_test
|
||||
env.curriculum.record_performance(generator_id_adv_test, True)
|
||||
|
||||
|
||||
# Try to advance difficulty
|
||||
did_advance = env.curriculum.advance_difficulty()
|
||||
new_level_manual_adv = env.curriculum.get_current_level()
|
||||
|
|
@ -204,9 +220,13 @@ async def main():
|
|||
logger.info(f" - Level before manual simulation: {initial_level_manual_adv}")
|
||||
logger.info(f" - Recorded {config.min_evaluations} correct answers manually.")
|
||||
logger.info(f" - Did advance: {did_advance}")
|
||||
logger.info(f" - Level after manual advancement attempt: {new_level_manual_adv}")
|
||||
logger.info(
|
||||
f" - Level after manual advancement attempt: {new_level_manual_adv}"
|
||||
)
|
||||
else:
|
||||
logger.info(f"Skipping manual advancement simulation as current level {initial_level_manual_adv} is already max level {max_level_possible}.")
|
||||
logger.info(
|
||||
f"Skipping manual advancement simulation as current level {initial_level_manual_adv} is already max level {max_level_possible}."
|
||||
)
|
||||
|
||||
logger.info("InfiniteMath local runner completed successfully.")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue