feat: Add iteration support to ArithmeticDataset with size-based termination

This commit is contained in:
Andreas Koepf (aider) 2025-01-23 12:34:54 +01:00
parent 3077d8571e
commit a8b3a86831
2 changed files with 39 additions and 0 deletions

View file

@ -141,6 +141,19 @@ class ArithmeticDataset:
expression = " ".join(expression_parts)
return expression, result
def __iter__(self):
"""Make the dataset iterable"""
self._current_idx = 0
return self
def __next__(self):
"""Get next item in iteration"""
if self._current_idx >= self.config.size:
raise StopIteration
item = self[self._current_idx]
self._current_idx += 1
return item
def _format_question(self, expression: str) -> str:
"""Format the expression according to config style"""
if self.config.format_style == "simple":