andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-10 00:28:31 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

feat: add ability to train on custom file (#1161)

Browse Source 
* feat: add ability to train on custom file

* feat: add pkl file validation

* feat: fix tests

* feat: fix tests

* feat: fix tests
This commit is contained in:
Eduardo Chiarotti
2024-08-09 19:41:58 -03:00
committed by GitHub
parent 62f5b2fb2e
commit 51ee483e9d
8 changed files with 54 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/crewai/cli/train_crew.py
View File

@@ -3,19 +3,22 @@ import subprocess
import click
def train_crew(n_iterations: int) -> None:
def train_crew(n_iterations: int, filename: str) -> None:
"""
Train the crew by running a command in the Poetry environment.
Args:
n_iterations (int): The number of iterations to train the crew.
"""
command = ["poetry", "run", "train", str(n_iterations)]
command = ["poetry", "run", "train", str(n_iterations), filename]
try:
if n_iterations <= 0:
raise ValueError("The number of iterations must be a positive integer.")
if not filename.endswith(".pkl"):
raise ValueError("The filename must not end with .pkl")
result = subprocess.run(command, capture_output=False, text=True, check=True)
if result.stderr: