add kaggle manager

This commit is contained in:
yzlin 2023-12-01 00:44:47 +08:00
parent 5522742002
commit 8b3d640dd6
3 changed files with 167 additions and 0 deletions

37
kaggle_team.py Normal file
View file

@ -0,0 +1,37 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import asyncio
import fire
from metagpt.roles.kaggle_manager import KaggleManager
from metagpt.roles.ml_engineer import MLEngineer
from metagpt.team import Team
async def main(
# competition: str,
# data_desc: str,
# requirement: str,
investment: float = 3.0,
n_round: int = 5,
):
competition, data_desc, requirement = (
"titanic",
"Training set is train.csv.\nTest set is test.csv. We also include gender_submission.csv, a set of predictions that assume all and only female passengers survive, as an example of what a submission file should look like.",
"Run EDA on the train dataset, train a model to predict survival (20% as validation) and save it, predict the test set using saved model, save the test result according to format",
)
team = Team()
team.hire(
[
KaggleManager(competition=competition, data_desc=data_desc),
MLEngineer(goal=requirement),
]
)
team.invest(investment)
team.start_project(requirement)
await team.run(n_round=n_round)
if __name__ == '__main__':
fire.Fire(main)