|
@@ -0,0 +1,156 @@
|
|
|
|
|
+import os
|
|
|
|
|
+import time
|
|
|
|
|
+import subprocess
|
|
|
|
|
+import sys
|
|
|
|
|
+import argparse
|
|
|
|
|
+import time
|
|
|
|
|
+from openai import OpenAI
|
|
|
|
|
+
|
|
|
|
|
+def fetch_generator(statement, model, maxtokens):
|
|
|
|
|
+ APIKEY = os.environ.get('DEEPSEEK_API_KEY')
|
|
|
|
|
+ if not APIKEY:
|
|
|
|
|
+ raise Exception('需要设置 DEEPSEEK_API_KEY 环境变量')
|
|
|
|
|
+
|
|
|
|
|
+ client = OpenAI(
|
|
|
|
|
+ api_key=APIKEY,
|
|
|
|
|
+ base_url='https://api.deepseek.com')
|
|
|
|
|
+
|
|
|
|
|
+ start_time = time.time()
|
|
|
|
|
+ msgs = [{ 'role': 'user', 'content':
|
|
|
|
|
+ """
|
|
|
|
|
+请根据以下信息学竞赛题目生成数据生成器。只输出 gen.py 代码。
|
|
|
|
|
+要求:
|
|
|
|
|
+1. 解析题目中的输入格式、数据范围和子任务划分
|
|
|
|
|
+2. 创建 data/ 目录并生成所有 .in 文件
|
|
|
|
|
+3. 数据严格符合题目约束
|
|
|
|
|
+4. 使用合理的数据梯度,包含边界情况
|
|
|
|
|
+5. 数据文件名一定为 'data{编号}.in',编号不含前导零
|
|
|
|
|
+6. 没有子任务的,数据文件直接放在 data/ 下;有子任务的,为每个子任务创建 subtask{编号}/ 目录,把数据文件移动到对应的子任务目录中
|
|
|
|
|
+7. 数据点在 10 到 20 个内
|
|
|
|
|
+8. 务必确保你给出的数据是正确的!
|
|
|
|
|
+输出只包含可以直接运行的 gen.py 的代码,无其他内容。
|
|
|
|
|
+题目描述:""" + statement}]
|
|
|
|
|
+ resp = client.chat.completions.create(
|
|
|
|
|
+ model=model,
|
|
|
|
|
+ messages=msgs,
|
|
|
|
|
+ stream=True,
|
|
|
|
|
+ temperature=0.0,
|
|
|
|
|
+ max_tokens=maxtokens)
|
|
|
|
|
+
|
|
|
|
|
+ response = ''
|
|
|
|
|
+ for c in resp:
|
|
|
|
|
+ d = c.choices[0].delta
|
|
|
|
|
+ if hasattr(d, 'reasoning_content') and d.reasoning_content is not None:
|
|
|
|
|
+ delta = d.reasoning_content
|
|
|
|
|
+ elif hasattr(d, 'content') and d.content is not None:
|
|
|
|
|
+ delta = d.content
|
|
|
|
|
+ if response == '':
|
|
|
|
|
+ delta = '\nThinking done\n' + delta
|
|
|
|
|
+ response += d.content
|
|
|
|
|
+ else:
|
|
|
|
|
+ continue
|
|
|
|
|
+
|
|
|
|
|
+ with open('stream', 'a') as f:
|
|
|
|
|
+ f.write(delta)
|
|
|
|
|
+ if hasattr(c, 'usage') and c.usage:
|
|
|
|
|
+ tokens = c.usage
|
|
|
|
|
+
|
|
|
|
|
+ if response.startswith('```python'):
|
|
|
|
|
+ response = response[9:-3]
|
|
|
|
|
+ return (response, tokens, time.time() - start_time)
|
|
|
|
|
+
|
|
|
|
|
+def execute(s):
|
|
|
|
|
+ print('正在执行命令:', s)
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ subprocess.run(['sh', '-c', s], check=True, timeout=20)
|
|
|
|
|
+ except subprocess.CalledProcessError as e:
|
|
|
|
|
+ raise Exception('命令执行错误')
|
|
|
|
|
+ except subprocess.TimeoutExpired:
|
|
|
|
|
+ raise Exception('命令执行超时(20s)')
|
|
|
|
|
+
|
|
|
|
|
+def generate_in():
|
|
|
|
|
+ # 其实应该用个沙盒比较安全,但是我懒得写了。可以最后扔到容器里头跑。
|
|
|
|
|
+ execute('python gen.py')
|
|
|
|
|
+
|
|
|
|
|
+def flatten_data():
|
|
|
|
|
+ execute('find -type f | xargs -I {} mv -n {} . || true')
|
|
|
|
|
+ execute('rmdir subtask* || true')
|
|
|
|
|
+
|
|
|
|
|
+def run_std():
|
|
|
|
|
+ if not os.path.exists('../std.cpp'):
|
|
|
|
|
+ raise Exception('未找到 std.cpp')
|
|
|
|
|
+ execute('g++ ../std.cpp -o std -Wall -Wextra -O2 -std=c++14')
|
|
|
|
|
+
|
|
|
|
|
+ for i in os.listdir():
|
|
|
|
|
+ if not i.endswith('.in'):
|
|
|
|
|
+ print(f'警告:未知文件 {i}')
|
|
|
|
|
+ continue
|
|
|
|
|
+ out = i[:-2] + 'ans'
|
|
|
|
|
+ execute(f'./std < {i} > {out}')
|
|
|
|
|
+ execute('rm std')
|
|
|
|
|
+
|
|
|
|
|
+def gen_problem_conf(time, mem):
|
|
|
|
|
+ execute(f'python ../autoconf.py {time} {mem}')
|
|
|
|
|
+
|
|
|
|
|
+def compress_data():
|
|
|
|
|
+ execute('zip ../data.zip -r .')
|
|
|
|
|
+
|
|
|
|
|
+def cleanup():
|
|
|
|
|
+ os.chdir('..')
|
|
|
|
|
+ execute('rm -r data/')
|
|
|
|
|
+
|
|
|
|
|
+def pretty_size(d):
|
|
|
|
|
+ if d < 1024:
|
|
|
|
|
+ return f'{d}'
|
|
|
|
|
+ if d < 1048576:
|
|
|
|
|
+ return f'{d/1024:.2f} KB'
|
|
|
|
|
+ return f'{d/1048576:.2f} MB'
|
|
|
|
|
+
|
|
|
|
|
+if __name__ == '__main__':
|
|
|
|
|
+ start_time = time.time()
|
|
|
|
|
+
|
|
|
|
|
+ parser = argparse.ArgumentParser()
|
|
|
|
|
+ parser.add_argument('statement', help='题面文件')
|
|
|
|
|
+ parser.add_argument('--model', type=str, default='deepseek-v4-pro', help='使用的模型(默认为 deepseek-v4-pro,可以改为 deepseek-v4-flash)')
|
|
|
|
|
+ parser.add_argument('--maxtokens', type=int, default=20000, help='token 上限(默认为 20000)')
|
|
|
|
|
+ parser.add_argument('--time', type=int, default=1, help='时间限制(秒,默认为 1)')
|
|
|
|
|
+ parser.add_argument('--mem', type=int, default=256, help='空间限制(MB,默认为 256)')
|
|
|
|
|
+ parser.add_argument('--offline', help='不生成新的 gen.py', action='store_true')
|
|
|
|
|
+ args = parser.parse_args()
|
|
|
|
|
+
|
|
|
|
|
+ if os.path.exists('data'):
|
|
|
|
|
+ print('正在清空已有的 data 目录')
|
|
|
|
|
+ execute('rm -r data')
|
|
|
|
|
+ if os.path.exists('data.zip'):
|
|
|
|
|
+ print('删除已存在的 data.zip')
|
|
|
|
|
+ execute('rm data.zip')
|
|
|
|
|
+
|
|
|
|
|
+ tokens = 0
|
|
|
|
|
+
|
|
|
|
|
+ with open(args.statement) as f:
|
|
|
|
|
+ if not args.offline:
|
|
|
|
|
+ print(f'正在请求 {args.model},token 上限为 {args.maxtokens} 个')
|
|
|
|
|
+ resp, tokens, duration = fetch_generator(f.read(), args.model, args.maxtokens)
|
|
|
|
|
+ with open('gen.py', 'w') as fw:
|
|
|
|
|
+ fw.write(resp)
|
|
|
|
|
+ print(f'--- tokens: {tokens}, time cost: {duration:.6f} seconds')
|
|
|
|
|
+ else:
|
|
|
|
|
+ if not os.path.exists('gen.py'):
|
|
|
|
|
+ raise Exception('离线模式下必须提供已有的 gen.py')
|
|
|
|
|
+ print('(离线模式)')
|
|
|
|
|
+
|
|
|
|
|
+ generate_in()
|
|
|
|
|
+ os.chdir('data')
|
|
|
|
|
+ gen_problem_conf(args.time, args.mem)
|
|
|
|
|
+
|
|
|
|
|
+ # run_std 需要一个扁平的目录,也就是所有文件放在 data/ 下,而 autoconf.py 不要求这么做
|
|
|
|
|
+ flatten_data()
|
|
|
|
|
+ run_std()
|
|
|
|
|
+ compress_data()
|
|
|
|
|
+ cleanup()
|
|
|
|
|
+ prettysize = pretty_size(os.path.getsize('data.zip'))
|
|
|
|
|
+ print(f'已生成 data.zip({prettysize}),耗时 {(time.time() - start_time):.2f} 秒,消耗 {tokens} 个 token')
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+
|