import%20marimo%0A%0A__generated_with%20%3D%20%220.17.6%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20ref%3A%20https%3A%2F%2Fdspy.ai%2Ftutorials%2Fgepa_aime%2F%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20random%0A%20%20%20%20from%20typing%20import%20Literal%0A%0A%20%20%20%20import%20dspy%0A%20%20%20%20from%20datasets%20import%20load_dataset%0A%20%20%20%20return%20Literal%2C%20dspy%2C%20load_dataset%2C%20random%0A%0A%0A%40app.cell%0Adef%20_(dspy)%3A%0A%20%20%20%20lm%20%3D%20dspy.LM(%22openai%2Fgpt-4.1-mini%22%2C%20temperature%3D1%2C%20max_tokens%3D32000)%0A%20%20%20%20dspy.configure(lm%3Dlm)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Loading%20the%20AIME%20dataset%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(dspy%2C%20load_dataset%2C%20random)%3A%0A%20%20%20%20def%20init_dataset()%3A%0A%20%20%20%20%20%20%20%20train_split%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20dspy.Example(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22problem%22%3A%20x%5B%22problem%22%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22solution%22%3A%20x%5B%22solution%22%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22answer%22%3A%20x%5B%22answer%22%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20).with_inputs(%22problem%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20x%20in%20load_dataset(%22AI-MO%2Faimo-validation-aime%22)%5B%22train%22%5D%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%20%20random.Random(0).shuffle(train_split)%0A%20%20%20%20%20%20%20%20total_num%20%3D%20len(train_split)%0A%20%20%20%20%20%20%20%20train_set%20%3D%20train_split%5B%3A%20int(0.5%20*%20total_num)%5D%0A%20%20%20%20%20%20%20%20val_set%20%3D%20train_split%5Bint(0.5%20*%20total_num)%3A%5D%0A%0A%20%20%20%20%20%20%20%20test_split%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20dspy.Example(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22problem%22%3A%20x%5B%22problem%22%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22answer%22%3A%20x%5B%22answer%22%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20).with_inputs(%22problem%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20x%20in%20load_dataset(%22MathArena%2Faime_2025%22)%5B%22train%22%5D%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%20%20test_set%20%3D%20test_split%20*%205%0A%0A%20%20%20%20%20%20%20%20return%20train_set%2C%20val_set%2C%20test_set%0A%20%20%20%20return%20(init_dataset%2C)%0A%0A%0A%40app.cell%0Adef%20_(init_dataset)%3A%0A%20%20%20%20train_set%2C%20val_set%2C%20test_set%20%3D%20init_dataset()%0A%0A%20%20%20%20len(train_set)%2C%20len(val_set)%2C%20len(test_set)%0A%20%20%20%20return%20test_set%2C%20train_set%0A%0A%0A%40app.cell%0Adef%20_(train_set)%3A%0A%20%20%20%20train_set%5B0%5D%5B%22problem%22%5D%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(train_set)%3A%0A%20%20%20%20train_set%5B0%5D%5B%22answer%22%5D%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(train_set)%3A%0A%20%20%20%20train_set%5B0%5D%5B%22solution%22%5D%5B%3A200%5D%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20A%20simple%20dspy.ChainOfThought%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(dspy)%3A%0A%20%20%20%20class%20GenerateResponse(dspy.Signature)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Solve%20the%20problem%20and%20provide%20the%20answer%20in%20the%20correct%20format.%22%22%22%0A%20%20%20%20%20%20%20%20problem%20%3D%20dspy.InputField()%0A%20%20%20%20%20%20%20%20answer%20%3D%20dspy.OutputField()%0A%20%20%20%20return%20(GenerateResponse%2C)%0A%0A%0A%40app.cell%0Adef%20_(GenerateResponse%2C%20dspy)%3A%0A%20%20%20%20program%20%3D%20dspy.ChainOfThought(GenerateResponse)%0A%20%20%20%20return%20(program%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Defining%20the%20evaluation%20metric%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Literal)%3A%0A%20%20%20%20def%20metric(example%2C%20prediction%2C%20trace%3DNone%2C%20pred_name%3DNone%2C%20pred_trace%3DNone)%20-%3E%20Literal%5B0%2C%201%5D%3A%0A%20%20%20%20%20%20%20%20correct_answer%20%3D%20int(example%5B%22answer%22%5D)%0A%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20llm_answer%20%3D%20int(prediction.answer)%0A%20%20%20%20%20%20%20%20except%20ValueError%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%200%0A%20%20%20%20%20%20%20%20return%20int(correct_answer%20%3D%3D%20llm_answer)%0A%20%20%20%20return%20(metric%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Evaluating%20unoptimized%20Chain%20Of%20Thought%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(dspy%2C%20metric%2C%20test_set)%3A%0A%20%20%20%20evaluate%20%3D%20dspy.Evaluate(%0A%20%20%20%20%20%20%20%20devset%3Dtest_set%2C%0A%20%20%20%20%20%20%20%20metric%3Dmetric%2C%0A%20%20%20%20%20%20%20%20num_threads%3D32%2C%0A%20%20%20%20%20%20%20%20display_table%3DTrue%2C%0A%20%20%20%20%20%20%20%20display_progress%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20return%20(evaluate%2C)%0A%0A%0A%40app.cell%0Adef%20_(evaluate%2C%20program)%3A%0A%20%20%20%20evaluate(program)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A
0e9996358a7be83c08d30daa2882dbdcaaf6bbaa274dbf18b9bb2c605c34f50f