## Math Agent distillation

This example demonstrates how to improve the math skills of a small LLama model. We will use the [GSM8k](https://huggingface.co/datasets/openai/gsm8k) dataset, which consists of 8,000 math word problems with their corresponding equations.

We built a basic math agent that uses the LLama 3.1 70B model equipped with reasoning and calculator tool to solve math problems: [math_agent.py](math_agent.py).

Steps to distill the math agent:

- [run it as a teacher](produce_teacher_tapes.py), collect the tapes of the successful solutions, and produce training data for the Math Agent from them. How to run: `uv run -m examples.gsm8k_tuning.produce_teacher_tapes` or download the data generated by `meta-llama/Llama-3.1-70B-Instruct`

```python
    df = pd.read_json("hf://datasets/ServiceNow/llama31_70b_gsm8k_3k/training_samples_3k.jsonl", lines=True)
    df.to_json("training_samples_3k.jsonl", orient="records", lines=True)
```

- [fine-tune smaller LLama 3.1 8B model](finetune_student.py) on the training data to get a tuned Math Agent. How to run: `uv run -m examples.gsm8k_tuning.finetune_student`
- [merge the lora weights](../../tapeagents/finetune/lora.py) to be able to serve the model with vLLM. How to run: `uv run -m tapeagents.finetune.lora PATH/TO/WEIGHTS`
- [evaluate the tuned Math Agent](evaluate_student.py) on the subset of GSM8K test set, comparing the accuracy of the teacher agent, student agent before tuning, and student agent after tuning. How to run: `uv run -m examples.gsm8k_tuning.evaluate_student`

<img width="526" alt="image" src="https://github.com/user-attachments/assets/55d099ab-ff5c-480b-b5b3-504b4206e677">

| Model | Test accuracy |
| ----- | ------------- |
| 8B student before tuning | 0.662 |
| 8B student after tuning | 0.785 |
| 70B teacher | 0.931 |

RL tuning on both successful and unsuccessful solutions is coming soon. Stay tuned!
