from dataloaders.dataset1d import EcgPipelineDataset1D
from models import models1d
from pipelines.base_pipeline import BasePipeline


class Pipeline1D(BasePipeline):
    def __init__(self, config):
        super().__init__(config)

    def _init_net(self):
        model = getattr(models1d, self.config["model"])(
            num_classes=self.config["num_classes"],
        )
        model = model.to(self.config["device"])
        return model

    def _init_dataloader(self):
        inference_loader = EcgPipelineDataset1D(self.config["ecg_data"]).get_dataloader(
            batch_size=self.config["batch_size"],
            num_workers=self.config["num_workers"],
            shuffle=False,
        )

        return inference_loader