from clearml import Task

class ClearMLHelper:
    def __init__(self, project_name: str):
        self.project_name = project_name

    def get_task(self, task_name: str = "Model Training"):
        Task.add_requirements("requirements.txt")
        Task.ignore_requirements("torch")
        Task.ignore_requirements("torchvision")
        Task.ignore_requirements("tensorboard")
        task = Task.init(project_name=self.project_name, task_name=task_name, continue_last_task=False)
        task.set_base_docker(f"docker.io/clearml/pytorch-cuda-gcc:2.0.0-cuda11.7-cudnn8-runtime")
        task.set_packages("requirements.txt")
        return task