AutoTabML / autotabml_tasks.py
singhtech's picture
Upload 3 files
715848f verified
from crewai import Task
# Function to create tasks based on user inputs
def create_tasks(func_call,user_question,file_name, data_upload, df, suggestion, edited_code, debugger, agents):
info = df.info()
tasks = []
if(func_call == "Process"):
tasks.append(Task(
description=f"Clarify the ML problem: {user_question}",
agent=agents["Problem_Definition_Agent"],
expected_output="A clear and concise definition of the ML problem."
)
)
if data_upload:
tasks.extend([
Task(
description=f"Evaluate the data provided by the file name . This is the data: {df}",
agent=agents["EDA_Agent"],
expected_output="An assessment of the EDA and preprocessing like dataset info, missing value, duplication, outliers etc. on the data provided"
),
Task(
description=f"Feature Engineering on data {df} based on EDA output: {info}",
agent=agents["Feature_Engineering_Agent"],
expected_output="An assessment of the Featuring Engineering and preprocessing like handling missing values, handling duplication, handling outliers, feature encoding, feature scaling etc. on the data provided"
)
])
tasks.extend([
Task(
description="Suggest suitable ML models.",
agent=agents["Model_Recommendation_Agent"],
expected_output="A list of suitable ML models."
),
Task(
description=f"Generate starter Python code based on feature engineering, where column names are {df.columns.tolist()}. Generate only the code without any extra text",
agent=agents["Starter_Code_Generator_Agent"],
expected_output="Starter Python code."
),
])
if(func_call == "Modify"):
if suggestion:
tasks.append(
Task(
description=f"Modify the already generated code {edited_code} according to the suggestion: {suggestion} \n\n Do not generate entire new code.",
agent=agents["Code_Modification_Agent"],
expected_output="Modified code."
)
)
if(func_call == "Debug"):
if debugger:
tasks.append(
Task(
description=f"Debug and fix any errors for data with column names {df.columns.tolist()} with data as {df} in the generated code: {edited_code} \n\n According to the debugging: {debugger}. \n\n Do not generate entire new code. Just remove the error in the code by modifying only necessary parts of the code.",
agent=agents["Code_Debugger_Agent"],
expected_output="Debugged and successfully executed code."
)
)
tasks.append(
Task(
description = "Your job is to only extract python code from string",
agent = agents["Compiler_Agent"],
expected_output = "Running python code."
)
)
return tasks