Spaces:
Running on Zero
Running on Zero
File size: 5,341 Bytes
0dd6c2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | import json
from types import TracebackType
from typing import Any
from linalg_zero.generator.models import ComponentResult, DifficultyCategory
from linalg_zero.generator.sympy.template_engine import TemplateEngine
from linalg_zero.shared.types import LibTypes
class ProblemContext:
"""
Context manager for state information around the resolution process.
"""
def __init__(self, entropy: float, difficulty_level: DifficultyCategory, step_counter: int):
self.entropy = entropy
self.difficulty_level = difficulty_level
self.used_entropy = 0.0
self.tool_calls_count = 0
self.stepwise_results: list[dict[str, Any]] = []
self.golden_result: dict[str, str] = {}
self._step_counter = step_counter
self.constraints: dict[str, Any] = {}
def __enter__(self) -> "ProblemContext":
return self
def __exit__(self, exc_type: type, exc_val: Exception, exc_tb: TracebackType) -> None:
pass
def record_entropy_usage(self, amount: float) -> None:
"""
Record entropy usage for tracking problem complexity.
"""
self.used_entropy += amount
def allocate_entropy(self, entropy: float | None) -> float:
"""
Resolve and consume an entropy amount based on the given value or use
entire entropy budget if None is provided. The provided value allows to
allocate entropy multiple times across a context lifetime.
The chosen amount is recorded against the context budget.
"""
remaining = self.entropy - self.used_entropy
if remaining <= 1e-12:
raise ValueError(f"Entropy budget exceeded: remaining {remaining:.3f}")
amount: float | None = None
if entropy is not None:
# If entropy is provided, use it directly.
# This can happen during generator lifetime that require multiple variable allocations.
amount = entropy
if amount is None:
amount = remaining
if amount > remaining + 1e-12:
raise ValueError(f"Entropy budget exceeded: request {amount:.3f}, remaining {remaining:.3f}")
self.record_entropy_usage(amount)
return amount
def _prepare_verification_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Prepare verification data by extracting dependencies and input fields."""
verification = {}
num_inputs = 0
num_dependencies = 0
# Handle dependencies
dependent_on = input_data.pop("dependent_on", None)
if dependent_on is not None:
verification["dependent_on"] = dependent_on
num_dependencies = len(dependent_on)
# Extract and JSON-encode all input_* fields
for key in list(input_data.keys()):
if key.startswith("input_"):
verification[key] = json.dumps(input_data.pop(key))
num_inputs += 1
assert num_inputs == num_dependencies, "Number of inputs and dependencies must match"
# Add generator type and remaining input data
verification["generator_type"] = input_data.pop("generator_type")
verification["input"] = json.dumps(input_data)
return verification
def record_tool_call(
self,
function_name: str,
result: LibTypes,
input_data: dict[str, Any],
is_final: bool = False,
) -> str:
"""
Record a tool call with its result. It tracks the dependencies between
steps which will later be used to verify correctness during GRPO.
"""
self.tool_calls_count += 1
step_id = str(self._step_counter)
if result is not None:
result_json = json.dumps(result)
step_data = {
"tool": function_name,
"result": result_json,
"step_id": step_id,
"verification": self._prepare_verification_data(input_data),
}
if is_final:
self.golden_result = {"final_answer": result_json, "from_step_id": step_id}
self.stepwise_results.append(step_data)
self._step_counter += 1
return step_id
class CompositionContext(ProblemContext):
"""
Extends the base ProblemContext to support shared state and global variables
across composed problem components.
"""
def __init__(
self,
entropy: float,
difficulty_level: DifficultyCategory,
step_counter: int,
template_engine: TemplateEngine,
local_index: int,
):
super().__init__(entropy, difficulty_level, step_counter)
self.component_results: list[ComponentResult] = []
self.template_engine = template_engine
self.local_index = local_index
def record_component_result(self, result: ComponentResult) -> None:
"""Record the result of a component execution."""
self.component_results.append(result)
# Update entropy usage and validate budget
self.used_entropy += result.entropy_consumed
if self.used_entropy > self.entropy + 1e-12:
raise ValueError(f"Entropy budget exceeded: used {self.used_entropy:.3f}, available {self.entropy:.3f}")
self.tool_calls_count += result.tool_calls_used
|