init-commit

internbootcamp/libs/game24/README.md

@@ -0,0 +1,141 @@
+# Game24Plus - 数学谜题生成与解决工具
+
+`Game24Plus` 是一个用于生成和解决类似于“24点游戏”的数学谜题的 Python 工具。它通过加、减、乘、除四种基本运算，将给定的数字组合成目标数字。该工具支持生成不同数量和范围的数字组合，并提供多种生成方式，包括单进程和多进程并行处理。
+
+---
+
+## 功能特性
+
+- **生成数学谜题**：生成指定数量和范围的数字组合，并通过随机运算生成目标数字。
+- **解决谜题**：递归地尝试所有可能的运算组合，找到达到目标数字的解决方案。
+- **多进程并行处理**：支持多进程并行生成谜题，提高生成效率。
+- **数据保存**：生成的谜题和解决方案以 JSONL 格式保存到文件中。
+
+
+## 代码结构
+
+### 主要类与方法
+
+#### `Game24Plus` 类
+
+- **`__init__(self, num_numbers, range_max, target_max, seed=None)`**  
+  初始化随机数生成器、数字数量、数字范围、目标数字范围等。
+
+- **`sample_one_number(self, num_min, num_max)`**  
+  生成一个指定范围内的随机整数。
+
+- **`get_numbers(self)`**  
+  生成一组随机整数并排序。
+
+- **`enumerate_all_numbers(self, num_numbers)`**  
+  枚举所有可能的数字组合（递归实现）。
+
+- **`sample_operation(self, numbers)`**  
+  从数字列表中随机选择两个数字和一个运算符，返回运算表达式。
+
+- **`calculate(self, n1, op, n2)`**  
+  根据运算符计算两个数字的结果，处理除法时使用分数。
+
+- **`get_target(self, numbers)`**  
+  通过随机运算将数字列表逐步减少，直到得到一个目标数字，并记录运算过程。
+
+- **`get_target_limit_range(self, numbers)`**  
+  在指定范围内生成目标数字，避免无效结果。
+
+- **`solve(self, numbers, target)`**  
+  递归地尝试所有可能的运算组合，找到达到目标数字的解决方案。
+
+### 主要函数
+
+- **`construct_game24_v1`**  
+  生成指定数量的谜题，并将结果保存为 JSONL 文件。
+
+- **`construct_helper`**  
+  辅助函数，用于生成指定目标范围内的谜题。
+
+- **`construct_game24_v3`**  
+  使用多进程并行生成谜题，提高生成效率。
+
+- **`construct_helper_for_v4`**  
+  辅助函数，用于多进程版本的谜题生成。
+
+- **`construct_game24_v4`**  
+  使用多进程和多队列实现高效的谜题生成。
+
+---
+
+## 使用示例
+
+### 生成谜题
+
+```python
+from game24plus import Game24Plus
+
+# 初始化
+game = Game24Plus(num_numbers=4, range_max=100, target_max=1000)
+
+# 生成一组数字
+numbers = game.get_numbers()
+print("生成的数字：", numbers)
+
+# 生成目标数字和运算过程
+target, operations = game.get_target_limit_range(numbers)
+print("目标数字：", target)
+print("运算过程：", operations)
+```
+
+### 解决谜题
+
+```python
+# 解决谜题
+solution = game.solve(numbers, target=24)
+if solution:
+    print("解决方案：", solution)
+else:
+    print("无解")
+```
+
+### 批量生成谜题
+
+```bash
+python game24plus.py --num_numbers 4 --output_dir data/gameX/raw_1104
+```
+
+---
+
+## 参数说明
+
+- **`num_numbers`**：生成谜题的数字数量（默认为 4）。
+- **`range_max`**：生成数字的最大值（默认为 101）。
+- **`target_max`**：目标数字的最大值（默认为 1000）。
+- **`num_samples`**：生成的谜题数量（默认为 200000）。
+- **`seed`**：随机数种子（默认为 1234）。
+- **`num_workers`**：多进程并行处理的工作进程数量（默认为 64）。
+- **`output_dir`**：输出文件的保存目录（默认为当前目录）。
+
+---
+
+## 输出文件格式
+
+生成的谜题和解决方案以 JSONL 格式保存，每行包含一个 JSON 对象，例如：
+
+```json
+{
+  "puzzle": "3 4 6 8",
+  "target": "24",
+  "operations": [
+    ["3", "*", "8", "24"],
+    ["4", "+", "6", "10"],
+    ["10", "-", "3", "7"]
+  ]
+}
+```
+
+---
+
+## 性能优化
+
+- **多进程并行处理**：通过 `multiprocessing` 模块实现高效的谜题生成。
+- **递归算法优化**：减少重复计算，提高解决谜题的效率。
+
+---

internbootcamp/libs/game24/game24.py

@@ -0,0 +1,336 @@
+import math
+import random
+import fractions
+import json
+import tqdm
+import os
+import numpy as np
+from multiprocessing import Pool, Queue, Process
+
+class Game24Plus:
+    """General version of game24 (Krypto), given N numbers less than X, find a way to get number M using 4 basic operations."""
+    def __init__(self, num_numbers, range_max, target_max, seed=None):
+        self.rng = np.random.default_rng(seed)
+        self.num_numbers = num_numbers
+        self.target_max = target_max
+        self.range_max = range_max
+        self.operations = ['+', '-', '*', '/']
+        
+    def sample_one_number(self, num_min, num_max):
+        return self.rng.integers(num_min, num_max, size=1)[0]
+        
+    def get_numbers(self):
+        numbers = self.rng.integers(0, self.range_max, size=self.num_numbers).tolist()
+        numbers.sort()
+        return numbers
+    
+    def enumerate_all_numbers(self, num_numbers):
+        if num_numbers == 1:
+            for i in range(1, self.range_max):
+                yield [i]
+            return
+        for i in range(1, self.range_max):
+            for j in self.enumerate_all_numbers(num_numbers - 1):
+                yield [i] + j            
+
+    def sample_operation(self, numbers: list):
+        n_pos = self.rng.choice(len(numbers), size=2, replace=False)
+        n1 = numbers[n_pos[0]]
+        n2 = numbers[n_pos[1]]
+        op = self.rng.choice(self.operations).item()
+        if op == '/':
+        #     if n2 == 0 or n1 % n2 != 0:
+            if n2 == 0:
+                return None
+        return n1, op, n2, n_pos[0], n_pos[1]
+    
+    def calculate(self, n1, op, n2):
+        if op == '+':
+            return n1 + n2
+        elif op == '-':
+            return n1 - n2
+        elif op == '*':
+            return n1 * n2
+        elif op == '/':
+            if n1 % n2 != 0:
+                return fractions.Fraction(n1, n2)
+            return n1 // n2
+    
+    def get_target(self, numbers: list):
+        current_numbers = numbers.copy()
+        current_operations = []
+        while len(current_numbers) > 1:
+            r = self.sample_operation(current_numbers)
+            if r is None:
+                continue
+            n1, op, n2, pos1, pos2 = r
+            current_numbers.pop(max(pos1, pos2))
+            current_numbers.pop(min(pos1, pos2))
+            res = self.calculate(n1, op, n2)
+            if isinstance(res, fractions.Fraction):
+                a, b = res.as_integer_ratio()
+                if a % b == 0:
+                    res = int(a/b)
+            current_operations.append((str(n1), op, str(n2), str(res)))
+            current_numbers.append(res)
+        return current_numbers[0], current_operations
+    
+    def get_target_limit_range(self, numbers):
+        target, operations = self.get_target(numbers)
+        num_try = 0
+        while not isinstance(target, int) or target < 0 or target > self.target_max:
+            target, operations = self.get_target(numbers)
+            num_try += 1
+            if num_try > 1000:
+                return None, None
+        return target, operations
+    
+    def solve(self, numbers, target):
+        def precise_calculate(n1, op, n2):
+            if op == '+':
+                return n1 + n2
+            elif op == '-':
+                return n1 - n2
+            elif op == '*':
+                return n1 * n2
+            elif op == '/':
+                if n1 % n2 != 0:
+                    return fractions.Fraction(n1, n2)
+                return n1 // n2
+            
+        if len(numbers) == 1:
+            if isinstance(numbers[0], int) and numbers[0] == target:
+                return []
+            else:
+                a, b = numbers[0].as_integer_ratio()
+                if a % b == 0 and int(a/b) == target:
+                    return []
+                return None
+        for i, n1 in enumerate(numbers):
+            for j, n2 in enumerate(numbers):
+                if i == j:
+                    continue
+                for op in self.operations:
+                    if op == '/':
+                        # if n2 == 0 or n1 % n2 != 0:
+                        if n2 == 0:
+                            continue
+                    res = precise_calculate(n1, op, n2)
+                    new_num = [n for k, n in enumerate(numbers) if k != i and k != j] + [res]
+                    new_how = [(str(n1), op, str(n2), str(res))]
+                    part_solution = self.solve(new_num, target)
+                    if part_solution is not None:
+                        return new_how + part_solution
+    def generate_ground_truth(self, numbers, target):
+        calc_steps = self.solve(numbers, target)
+        
+        result_map = {}
+
+        for step in calc_steps:
+            operand1, operator, operand2, result = step
+            
+            # 如果操作数是之前的结果，则替换为对应的表达式
+            if operand1 in result_map:
+                operand1 = f"({result_map[operand1]})"
+            else:
+                operand1 = operand1
+            
+            if operand2 in result_map:
+                operand2 = f"({result_map[operand2]})"
+            else:
+                operand2 = operand2
+            
+            # 构建当前步骤的表达式
+            if operator == "":
+                expression = operand1  # 如果没有操作符，直接取操作数
+            else:
+                expression = f"{operand1} {operator} {operand2}"
+            
+            # 将当前结果存入映射表
+            result_map[result] = expression
+
+        # 最终结果是最后一个元组的结果
+        final_result = calc_steps[-1][-1]
+        final_expression = result_map[final_result]
+        
+        return final_expression
+            
+                    
+def construct_game24_v1(num_numbers=3, range_max=101, num_samples=10000, target_max=1000, seed=1234, output_dir=None):
+    game = Game24Plus(num_numbers, range_max, target_max=target_max, seed=seed)
+    with open(os.path.join(output_dir, f'train_m={num_numbers}.jsonl'), 'w') as fp:
+        puzzle_dict = {}
+        for idx in tqdm.trange(num_samples):
+            num_try = 0
+            while True:
+                numbers = game.get_numbers()
+                puzzle = ' '.join(str(n) for n in numbers)
+                if puzzle in puzzle_dict:
+                    num_try += 1
+                    if num_try > 1000:
+                        puzzle = None
+                        break
+                    continue
+                else:
+                    break
+            if puzzle is None:
+                print('Failed to generate unique puzzle')
+                break
+            puzzle_dict[puzzle] = True
+            num_try = 0
+            data = {}
+            while True:
+                target, operations = game.get_target_limit_range(numbers)
+                target_str = str(target)
+                key = puzzle + ' ' + target_str
+                if key in data:
+                    num_try += 1
+                    if num_try > 10:
+                        break
+                else:
+                    data[key] = {
+                        'puzzle': puzzle,
+                        'target': target_str,
+                        'operations': operations
+                    }
+                    num_try = 0
+            for key, ex in data.items():
+                print(json.dumps(ex), file=fp)
+
+
+def construct_helper(args):
+    num_numbers, range_max, num_samples_per_target, target_min, target_max, seed = args
+    game = Game24Plus(num_numbers, range_max, target_max=target_max, seed=seed)
+    data = []
+    for target in range(target_min, target_max):
+        num_try = 0
+        target_str = str(target)
+        seen_data = {}
+        while True:
+            numbers = game.get_numbers()
+            puzzle = ' '.join(str(n) for n in numbers)
+            if puzzle in seen_data:
+                continue
+            operations = game.solve(numbers, target)
+            if operations is None:
+                num_try += 1
+                if num_try > 10000:
+                    break
+                continue
+            ex = {
+                'puzzle': puzzle,
+                'target': target_str,
+                'operations': operations
+            }
+            data.append(ex)
+            num_try = 0
+            seen_data[puzzle] = 1
+            if len(seen_data) > num_samples_per_target:
+                break
+            
+    return data
+    
+
+def construct_game24_v3(num_numbers=3, range_max=101, num_samples=50000, target_max=1000, seed=1234, num_workers=64, output_dir=None):
+    print(f'Args', num_numbers, range_max, num_samples, target_max, seed, num_workers)
+    arg_list = []
+    num_samples_per_target = (num_samples // target_max) + 1
+    chunk_size = 1
+    start_target = 0
+    while start_target < target_max:
+        _target_min = start_target
+        _target_max = min(target_max, start_target + chunk_size)
+        arg_list.append((
+            num_numbers, range_max, num_samples_per_target, _target_min, _target_max, seed
+        ))
+        start_target = _target_max
+    # print(arg_list)
+    with open(os.path.join(output_dir, f'train_m={num_numbers}.jsonl'), 'w') as fp:
+        with Pool(num_workers) as p:
+            for data in tqdm.tqdm(p.imap_unordered(construct_helper, arg_list), total=len(arg_list)):
+                for ex in data:
+                    print(json.dumps(ex), file=fp)
+                    
+                    
+def construct_helper_for_v4(args_queue, out_queue):
+    while True:
+        args = args_queue.get()
+        if args is None:
+            break
+        numbers = args[0]
+        game = args[1]
+        puzzle = ' '.join(str(n) for n in numbers)
+        num_try = 0
+        data = {}
+        while True:
+            target, operations = game.get_target_limit_range(numbers)
+            if target in data:
+                num_try += 1
+                if num_try > 10:
+                    break
+            else:
+                data[target] = {
+                    'puzzle': puzzle,
+                    'target': str(target),
+                    'operations': operations
+                }
+        out_queue.put(data)
+    
+
+def construct_game24_v4(num_numbers=3, range_max=101, num_samples=200000, target_max=1000, seed=1234, num_workers=64, output_dir=None):
+    print(f'Args', num_numbers, range_max, num_samples, target_max, seed, num_workers)
+    seen_numbers = set()
+    procs = []
+    args_queue = Queue()
+    out_queue = Queue()
+    for i in range(num_workers):
+        p = Process(target=construct_helper_for_v4, args=(args_queue, out_queue))
+        p.start()
+        procs.append(p)
+        
+    game = Game24Plus(num_numbers, range_max, target_max=target_max, seed=seed)
+    with open(os.path.join(output_dir, f'train_m={num_numbers}.jsonl'), 'w') as fp:
+        for i in range(num_samples):
+            numbers = game.get_numbers()
+            puzzle = ' '.join(str(n) for n in numbers)
+            if puzzle in seen_numbers:
+                continue
+            seen_numbers.add(puzzle)
+            args_queue.put((numbers, game))
+        for i in range(num_workers):
+            args_queue.put(None)
+            
+        count = 0
+        tqdm_bar = tqdm.tqdm(total=num_samples)
+        while count < num_samples:
+            data = out_queue.get()
+            for target, ex in data.items():
+                print(json.dumps(ex), file=fp)
+                count += 1
+                tqdm_bar.update(1)
+                
+        for p in procs:
+            p.kill()
+            
+
+if __name__ == "__main__":
+    # 初始化 Game24Plus 实例
+    game = Game24Plus(num_numbers=6, range_max=10, target_max=24, seed=46000767)
+    
+    # 获取一组随机数字
+    numbers = game.get_numbers()
+    print("Numbers:", numbers)
+    
+    # 设置目标值
+    target = 24
+    print("Target:", target)
+    
+    print("Trace:", game.solve(numbers, target))
+    
+    # 尝试生成一步答案算式
+    expression = game.generate_ground_truth(numbers, target)
+    if expression:
+        print("expression:", expression)
+    else:
+        print("No expression found.")
+    

internbootcamp/libs/game24/main.py

@@ -0,0 +1,32 @@
+import random
+import numpy as np
+import fractions
+from game24 import Game24Plus
+
+
+# 实例化 Game24Plus
+game = Game24Plus(num_numbers=4, range_max=10, target_max=24, seed=random.randint(0, 1000))
+
+# 生成一组数字
+numbers = game.get_numbers()
+print("生成的数字:", numbers)
+
+# 计算目标值及其运算过程
+target, operations = game.get_target_limit_range(numbers)
+if target is not None:
+    print("计算得到的目标值:", target)
+    print("运算过程:")
+    for step in operations:
+        print(f"{step[0]} {step[1]} {step[2]} = {step[3]}")
+else:
+    print("未能生成符合条件的目标值。")
+
+# 求解问题：找到一种运算方式得到目标值
+if target is not None:
+    solution = game.solve(numbers, target)
+    if solution is not None:
+        print("求解得到的运算步骤:")
+        for step in solution:
+            print(f"{step[0]} {step[1]} {step[2]} = {step[3]}")
+    else:
+        print("未能找到符合条件的运算步骤。")