Sweep: refactor the training loop from a script to a class in main.py

Here's the PR! https://github.com/sweepai/evals/pull/163.

💎 Sweep Pro: I'm using GPT-4. You have unlimited GPT-4 tickets. (tracking ID: c863ad18ba)

Actions (click)

[ ] ↻ Restart Sweep

Sandbox Execution ✓

Here are the sandbox execution logs prior to making any changes:

Sandbox logs for 7284908

trunk fmt src/main.py || exit 0 1/2 ✓

  NOTICES  
 src/main.py  ignored empty file [black]
 Hint: use --force to check ignored files
Checked 1 file
✔ No issues

trunk check --fix --filter=-ruff --print-failures src/main.py 2/2 ✓

  NOTICES  
 src/main.py  ignored empty file [black]
 Hint: use --force to check ignored files
Checked 1 file
✔ No issues
Run trunk upgrade to upgrade 1 linter

Sandbox passed on the latest main, so sandbox checks will be enabled for this issue.

Step 1: 🔎 Searching

I found the following snippets in your repository. I will now analyze these snippets and come up with a plan.

Some code snippets I looked at (click to expand). If some file is missing from here, you can mention the path in the ticket description.

https://github.com/sweepai/evals/blob/72849082ab1fcf90323a2d3b20d400de0d41ebd9/src/main.py#L1-L46 https://github.com/sweepai/evals/blob/72849082ab1fcf90323a2d3b20d400de0d41ebd9/src/api.py#L1-L27 https://github.com/sweepai/evals/blob/72849082ab1fcf90323a2d3b20d400de0d41ebd9/requirements.txt#L1-L15

Step 2: ⌨️ Coding

[X] Modify src/main.py ✓ https://github.com/sweepai/evals/commit/35f37cdbd49a63bfdd7a39e2d65b0c186ad83c49

Modify src/main.py with contents:
• Create a new class named `Trainer` at the top of the file.
• Move the model, optimizer, and loss function initialization from the global scope into the `__init__` method of the `Trainer` class. The `__init__` method should take the learning rate and the path to save the model as parameters, and use them to initialize the optimizer and save the model respectively.
• Move the training loop from the global scope into a new method of the `Trainer` class named `train`. This method should take the number of epochs as a parameter.
• Move the model saving code from the global scope into a new method of the `Trainer` class named `save_model`.
• Remove the global code for model initialization, training, and saving, as it is now encapsulated in the `Trainer` class.
• At the end of the file, create an instance of the `Trainer` class, and call its `train` and `save_model` methods to train and save the model.
[X] Ran sandbox for src/main.py. ✓ https://github.com/sweepai/evals/commit/52cb808b1b8bc91787d586dd565b782319ec01c4

Sandbox logs for https://github.com/sweepai/evals/commit/52cb808b1b8bc91787d586dd565b782319ec01c4

trunk fmt src/main.py || exit 0 1/2 ✓
 ✔ Formatted src/main.py
Re-checking autofixed files...

Checked 1 file
✔ No issues
Run trunk upgrade to upgrade 1 linter
trunk check --fix --filter=-ruff --print-failures src/main.py 2/2 ✓
Checked 1 file
✔ No issues

[X] Modify src/api.py ✓ https://github.com/sweepai/evals/commit/c1333ece180d5e2412e2ec2e2d68d752b38fbdc0

Modify src/api.py with contents:
• Import the `Trainer` class from main.py at the top of the file.
• Replace the global model loading code with the creation of a `Trainer` instance and a call to its `load_model` method.
• Replace the global model usage in the `predict` function with a call to a method of the `Trainer` instance that returns the model.
[X] Ran sandbox for src/api.py. ✓ https://github.com/sweepai/evals/commit/b31fc724747dc0cd6fdae265be435aa36b06ab4a

Sandbox logs for https://github.com/sweepai/evals/commit/b31fc724747dc0cd6fdae265be435aa36b06ab4a

trunk fmt src/api.py || exit 0 1/2 ✓
 ✔ Formatted src/api.py
Re-checking autofixed files...

 ✔ Formatted src/api.py
Re-checking autofixed files...

Checked 1 file
✔ No issues
Run trunk upgrade to upgrade 1 linter
trunk check --fix --filter=-ruff --print-failures src/api.py 2/2 ✓
Checked 1 file
✔ No issues

Step 3: 🔁 Code Review

I have finished reviewing the code for completeness. I did not find errors for sweep/refactor-training-loop_27.

🎉 Latest improvements to Sweep:

Sweep now uses the rope library to refactor Python! Check out Large Language Models are Bad at Refactoring Code. To have Sweep refactor your code, try sweep: Refactor <your_file>.py!
Sweep finds and fixes bugs in your code by writing unit tests! Check out Having GPT-4 Iterate on Unit Tests like a Human.

💡 To recreate the pull request edit the issue title or description. To tweak the pull request, leave a comment on the pull request. ^{Join Our Discord}

sweepai / evals