# This is the code I want to move out
def do_something_with_data(data: pd.DataFrame):
   # Some complex logic
   return data


@task
def get_data() -> pd.DataFrame:
    """Get the wine dataset."""
    return load_wine(as_frame=True).frame

@task
def process_data(data: pd.DataFrame) -> pd.DataFrame:
    new_data = do_something_with_data(data)
    return new_data

@task
def train_model(data: pd.DataFrame, hyperparameters: dict) -> LogisticRegression:
   #model.fit()
   return model

@workflow
def training_workflow(hyperparameters: dict) -> LogisticRegression:
    data = get_data()
    processed_data = process_data(data=data)
    return train_model(
        data=processed_data,
        hyperparameters=hyperparameters,
    )

try:
   from module import do_something_with_data
except ImportError:
   from .module import do_something_with_data

@task
def get_data() -> pd.DataFrame:
    """Get the wine dataset."""
    return load_wine(as_frame=True).frame

@task
def process_data(data: pd.DataFrame) -> pd.DataFrame:
    new_data = do_something_with_data(data)
    return new_data

@task
def train_model(data: pd.DataFrame, hyperparameters: dict) -> LogisticRegression:
   #model.fit()
   return model

@workflow
def training_workflow(hyperparameters: dict) -> LogisticRegression:
    data = get_data()
    processed_data = process_data(data=data)
    return train_model(
        data=processed_data,
        hyperparameters=hyperparameters,
    )