RP3Net 0.0.1__py3-none-any.whl

RP3Net/training/cli.py

@@ -0,0 +1,166 @@
+import re
+import os
+import socket
+import sys
+import logging
+import lightning.pytorch as L
+import lightning.pytorch.utilities as L_util
+import lightning.pytorch.cli as L_cli
+import lightning.pytorch.callbacks as L_cb
+import lightning.pytorch.loggers as L_log
+import wandb
+
+from . import lm
+from .. import util
+
+log = util.get_logger(__name__)
+
+def setup_logging_torch(args):
+    log_level = args.log_level
+    os.environ["PP_LOG_LEVEL"] = log_level
+    logfile = args.logfile
+    if logfile is not None:
+        logfile_base = util.resolve(logfile)
+        logfile_base = str(logfile_base.parent/logfile_base.stem)
+        os.environ["PP_LOGFILE_BASE"] = logfile_base
+        # /homes/evgeny/micromamba/envs/ai/lib/python3.11/site-packages/lightning/fabric/plugins/environments/slurm.py
+        # SLURMEnvironment.world_size() and SLURMEnvironment.global_rank()
+        if int(os.environ.get("SLURM_NTASKS", "0")) > 1:
+            logfile = logfile.replace(".log", f"_{os.environ.get('SLURM_PROCID', '0')}.log")
+    util.setup_logging(logfile, log_level, log_console=logfile is None)
+    ll = logging.getLogger("lightning")
+    ll.propagate = True
+    ll.handlers.clear()
+    ll = logging.getLogger("lightning.pytorch")
+    ll.handlers.clear()
+    ll.propagate = True
+    log.info(f"Host: {socket.gethostname()}; PID: {os.getpid()}; Command line: {' '.join(sys.argv)}")
+
+class RP3Cli(L_cli.LightningCLI):
+    def __init__(self, *args, **kwargs):
+        self.wandb_logger = None
+        super().__init__(*args, **{**kwargs, 'save_config_kwargs': {"overwrite": True}})
+
+    @staticmethod
+    @L_util.rank_zero_only
+    def wandb_init(wandb_project, wandb_run_name, wandb_run_id):
+        return wandb.init(id=wandb_run_id, project=wandb_project, name=wandb_run_name, resume='allow')
+
+    def wandb_logger_init(self, config):
+        if 'wandb' not in config or config.wandb is None or \
+            'project' not in config.wandb or ('run' not in config.wandb and 'run_id' not in config.wandb) or \
+                config.wandb.project is None or (config.wandb.run is None and config.wandb.run_id is None) or \
+                ('disable' in config.wandb and config.wandb.disable):
+            log.info("No wandb logging")
+            return
+        else:
+            log.info("Configure wandb logging")
+        run_id = config.wandb.run_id if 'run_id' in config.wandb and config.wandb.run_id is not None \
+            else config.wandb.run
+        run = self.wandb_init(wandb_project=config.wandb.project, wandb_run_name=config.wandb.run, wandb_run_id=run_id)
+        if run is not None:
+            log.info(f"Wandb run: {run.name}({run.id})")
+        logger = L_log.WandbLogger(project=config.wandb.project, name=config.wandb.run, id=run_id)
+        return logger
+
+    def add_arguments_to_parser(self, parser: L_cli.LightningArgumentParser) -> None:
+        parser.add_argument("--logfile", help="Log file. Log output to console if set to None.")
+        parser.add_argument("--log_level", default="info",
+                            help="Log level of root logger. Appender levels are appropriately hard coded.")
+        parser.add_argument("--track_metric_checkpoints", choices=["last", "all", "best"],
+                            help="""
+Track checkpoints for training and validation metrics from the module. 
+If not provided, no checkpoints will be recorded at all.
+If any value is provided, only the best checkpoint will be recorded for all the metrics specified by `model.metrics_for_checkpointing()`.
+The value of this argument affects how checkpoints will be saved for 'train_loss'.
+                            """)
+        parser.add_argument("--wandb.project", help="Wandb project name", default=None)
+        parser.add_argument("--wandb.run", help="Wandb run name", default=None)
+        parser.add_argument("--wandb.run_id", help="Wandb run id, same as name by default", default=None)
+        parser.add_argument("--wandb.disable", help="Set to true to turn off wandb logging, without removing the rest of wandb settings", action='store_true')
+        parser.add_argument("--test_after_fit_metric", help="Metric to use for test_after_fit. If not set, then do not run test_after_fit", default=None)
+        parser.add_argument("--emlc_k", help="Number of student gradinent steps to perform per teacher step for EMLC", default=1, type=int)
+
+    def before_instantiate_classes(self) -> None:
+        config = self.config.get(str(self.subcommand), self.config)
+        setup_logging_torch(config)
+        self.wandb_logger = self.wandb_logger_init(config)
+
+    
+    def instantiate_trainer(self, **kwargs) -> L.Trainer:
+        log.info("Instantiating trainer")
+        config = self.config.get(str(self.subcommand), self.config)
+        metric_checkpoints = self._get(self.config_init, 'track_metric_checkpoints')
+        # metric_checkpoints = bool(metric_checkpoints)
+        if metric_checkpoints is not None:
+            self.init_metric_checkpoints(metric_checkpoints)
+        self.add_loggers(config)
+        return super().instantiate_trainer(**kwargs)
+
+    def add_loggers(self, config):
+        configured_loggers = self._get(self.config_init, 'trainer.logger', default=[])
+        if configured_loggers == True or configured_loggers is None:
+            configured_loggers = []
+        elif isinstance(configured_loggers, L_log.Logger):
+            configured_loggers = [configured_loggers]
+        elif configured_loggers == False:
+            return
+        add_csv_logger = True
+        for logger in configured_loggers:
+            if isinstance(logger, L_log.CSVLogger):
+                add_csv_logger = False
+                break
+        if add_csv_logger:
+            csv_logger = L_log.CSVLogger(config['trainer']['default_root_dir'])
+            configured_loggers.append(csv_logger)
+        if self.wandb_logger is not None:
+            configured_loggers.append(self.wandb_logger)
+        config_init = self.config_init.get(str(self.subcommand), self.config_init)
+        config_init['trainer']['logger'] = configured_loggers
+
+    def init_metric_checkpoints(self, checkpoint_save_flag):
+        model: lm.RP3LM = self.model
+        metrics = model.metrics.metrics_for_checkpointing()
+        default_root_dir = util.resolve(self._get(self.config, 'trainer.default_root_dir'))
+        if metrics is None:
+            return
+        trainer_config = self._get(self.config_init, 'trainer', default={})
+        if 'callbacks' not in trainer_config or trainer_config['callbacks'] is None:
+            callbacks = []
+            trainer_config['callbacks'] = callbacks
+        else:
+            callbacks = trainer_config['callbacks']
+            if isinstance(callbacks, L.Callback):
+                callbacks = [callbacks]
+                trainer_config['callbacks'] = callbacks
+        for i, (key, metric) in enumerate(metrics.items()):
+            mode = 'max' if metric.higher_is_better else 'min'
+            callbacks.append(L_cb.ModelCheckpoint(
+                dirpath=default_root_dir, monitor=key, mode=mode,
+                filename='{epoch}_{'+key+':.2f}',
+                save_on_train_epoch_end=False, save_weights_only=False
+            ))
+        assert checkpoint_save_flag in ['last', 'all', 'best']
+        if checkpoint_save_flag == 'best':
+            save_top_k = 1
+        elif checkpoint_save_flag == 'all':
+            save_top_k = -1
+        else:
+            save_top_k = 0
+        callbacks.append(L_cb.ModelCheckpoint(dirpath=default_root_dir, monitor='train_loss', filename='{epoch}_{train_loss:.2f}', mode='min',
+                                                      save_top_k=save_top_k, save_last=True))
+
+    def test_after_fit(self, metric):
+        model: lm.RP3LM = self.model
+        dm = self.datamodule
+        dir = util.resolve(self.trainer.default_root_dir)
+        filename_pattern = re.compile(r'^epoch=\d+_' + metric + r'=\d+(\.\d+)?\.ckpt$')
+        cp_file = util.find_checkpoint_file(dir, filename_pattern)
+        log.info(f"Loading checkpoint {cp_file}")
+        self.trainer.test(model, dm, ckpt_path=str(cp_file))
+
+    def after_fit(self):
+        metric = self._get(self.config, 'test_after_fit_metric')
+        if metric is not None:
+            log.info(f"Running test_after_fit on {metric}")
+            self.test_after_fit(metric)

RP3Net/training/data.py

@@ -0,0 +1,300 @@
+import os
+import typing
+import functools
+import zipfile
+
+import torch
+import torch.utils.data as torch_data
+import polars as pl
+import numpy as np
+import lightning as L
+import ml_collections as mlc
+
+from .lm import RP3LM
+from .. import util
+from .. import model
+
+log = util.get_logger(__name__)
+
+FULL_DF_DTYPE_PL = {'created_at': pl.Datetime(), 'source': pl.Categorical(), 'sub_source': pl.Categorical(),
+                 'no_tags_cluster_40_id': pl.String(), 'with_tags_cluster_90_id': pl.String(),
+                 'has_dna': pl.Boolean(),
+                 'experiment_id': pl.String(), 
+                 'yield_binary': pl.Boolean(), 'yield_cat': pl.Int64(),
+                 'host': pl.Categorical(), 'exp_outcome': pl.Categorical(), 
+                 'id': pl.String(), 'fasta_id': pl.String(),  'dna_fasta_id': pl.String(), 'fasta_id_no_tags': pl.String(),
+                 'ds_type': pl.Categorical(),
+                 'n_tags_end': pl.Int64(), 'c_tags_start': pl.Int64(), 'n_fragments': pl.Int64(), 'unique_target_count': pl.Int64(),
+                 'fasta_id_no_tags': pl.String(),
+                 'uniprot_id': pl.String(), 'gene_id': pl.String(), 'taxon_id': pl.Int64(),
+                 }
+
+def read_full_df_pl(path:str|os.PathLike, **kwargs) -> pl.DataFrame:
+    _schema = FULL_DF_DTYPE_PL if 'schema_overrides' not in kwargs else kwargs['schema_overrides']
+    kwargs.pop('schema_overrides', None)
+    df = pl.read_csv(path, schema_overrides=_schema, **kwargs)
+    return df
+
+def load_global_embeddings_file(embeddings_file:os.PathLike) -> typing.Mapping:
+    log.info(f"Loading global embeddings from {embeddings_file}")
+    embeddings_data = torch.load(embeddings_file)
+    ids = embeddings_data['ids']
+    embeddings = embeddings_data['embeddings']
+    return {id: embeddings[i] for i, id in enumerate(ids)}
+
+class RP3GlobalEmbeddingsDataSet(torch_data.Dataset):
+    def __init__(self, df: pl.DataFrame, prefix:str, embeddings:typing.Mapping[str, torch.Tensor]) -> None:
+        super().__init__()
+        self.df = df
+        self.prefix=prefix
+        self.embeddings = embeddings
+
+    def __len__(self):
+        return self.df.shape[0]
+    
+    def __getitem__(self, idx):
+        try:
+            row = self.df.row(idx, named=True)
+            ret = {
+                'idx': row['src_idx'], 
+                'yield_binary': int(row['yield_binary']),
+                'source': row['source'],
+                'embeddings': self.embeddings[row['id']]
+            }
+            log.debug(f"{self.prefix}: (torch={idx}, csv={ret['idx']}){row['id']}: {ret['yield_binary']}")
+            return ret
+        except Exception as e:
+            log.error(f"Top level catch in {self.prefix} __getitem__", exc_info=e)
+            raise e
+
+
+class RP3SequenceEmbeddingsDataSet(torch_data.Dataset):
+    def __init__(self, df: pl.DataFrame, prefix:str, embeddings:zipfile.ZipFile) -> None:
+        super().__init__()
+        self.df = df
+        self.prefix=prefix
+        self.embeddings:zipfile.ZipFile = embeddings
+
+    def __len__(self):
+        return self.df.shape[0]
+    
+    def __getitem__(self, idx):
+        try:
+            row = self.df.row(idx, named=True)
+            with self.embeddings.open(row['id'], 'r') as f:
+                emb = torch.load(f, weights_only=True)
+            ret = {
+                'idx': row['src_idx'], 
+                'yield_binary': int(row['yield_binary']),
+                'source': row['source'],
+                'embeddings': emb
+            }
+            log.debug(f"{self.prefix}: (torch={idx}, csv={ret['idx']}){row['id']}: {ret['yield_binary']}")
+            return ret
+        except Exception as e:
+            log.error(f"Top level catch in {self.prefix} __getitem__", exc_info=e)
+            raise e
+
+class RP3SequenceDataSet(torch_data.Dataset):
+    def __init__(self, df: pl.DataFrame, prefix, rng:np.random.Generator, max_seq_len:int=0):
+        super().__init__()
+        self.rng = rng
+        self.df = df
+        self.prefix=prefix
+        self.max_seq_len = max_seq_len
+
+    def seq_chunk(self, seq:str):
+        if self.max_seq_len == 0 or len(seq) <= self.max_seq_len:
+            log.debug(f"Not changing sequence of length {len(seq)}; max_seq_len={self.max_seq_len}")
+            return seq
+        start_idx = self.rng.integers(len(seq) - self.max_seq_len + 1)
+        end_idx = start_idx + self.max_seq_len
+        log.debug(f"Returning the {start_idx}:{end_idx} chunk from sequence of length {len(seq)}; max_seq_len={self.max_seq_len}")
+        return seq[start_idx:end_idx]
+    
+    def __len__(self):
+        return self.df.shape[0]
+    
+    def __getitem__(self, idx):
+        try:
+            row = self.df.row(idx, named=True)
+            ret = {
+                'idx': row['src_idx'],
+                'source': row['source'],
+                'seq': self.seq_chunk(row['seq']), 
+                'yield_binary': int(row['yield_binary']),
+            }
+            log.debug(f"{self.prefix}: (torch={idx}, csv={ret['idx']}){row['id']}: {ret['yield_binary']}")
+            return ret
+        except Exception as e:
+            log.error(f"Top level catch in {self.prefix} __getitem__", exc_info=e)
+            raise e
+
+
+class RP3LDM(L.LightningDataModule):
+    def __init__(self, hypers) -> None:
+        super().__init__()
+        log.debug("DataModule init")
+        self.save_hyperparameters({'data': hypers})
+        self.hypers = mlc.ConfigDict(self.hparams.data)
+        self.sources_map = {s: i for i, s in enumerate(self.hypers.sources)}
+        self.rng = np.random.default_rng(self.hypers.get('seed', None))
+        self.validation_slice = self.hypers.get('validation_slice', 'VALIDATION')
+     
+    def torch_dataset(self, df:pl.DataFrame, prefix:str) -> torch_data.Dataset:
+        raise NotImplemented()
+
+    def load_df(self) -> pl.DataFrame:
+        data_path = self.hypers.ds_path
+        log.info(f"Loading data from {data_path}; validation slice: {self.validation_slice}")
+        df = read_full_df_pl(data_path).with_row_index('src_idx')
+        df_sources = set(*df.select(pl.col('source').cast(pl.String).unique()))
+        for s in self.sources_map:
+            assert s in df_sources
+        df = (df
+            .filter(pl.col('source').is_in(self.hypers.sources))
+            .with_columns(pl.col('source').cast(pl.String).replace_strict(self.sources_map))
+        )
+        return df
+    
+    def setup(self, stage: str) -> None:
+        log.debug("RP3LDM setup")
+        if self.trainer is not None:
+            assert self.trainer.model.sources == self.hypers.sources
+        df = self.load_df()
+        self.df_train = df.filter(
+            pl.col('ds_type').is_not_null() &
+            pl.col('ds_type').is_in(['TEST', self.validation_slice]).not_()
+        )
+        assert self.df_train.shape[0] > 0, f"No training data for slice {self.validation_slice}"
+        self.df_val = df.filter(ds_type=self.validation_slice)
+        assert self.df_val.shape[0] > 0, f"No validation data for slice {self.validation_slice}"
+        self.df_test = df.filter(ds_type='TEST')
+        assert self.df_test.shape[0] > 0, f"No test data for slice {self.validation_slice}"
+        self.create_torch_datasets()
+
+    def create_torch_datasets(self):
+        self.train_ds = self.torch_dataset(self.df_train, 'train')
+        self.val_ds = self.torch_dataset(self.df_val, "val")
+        df_val_train = self.df_train.sample(len(self.val_ds), with_replacement=False, seed=self.hypers.test_val_seed)
+        self.val_train_ds = self.torch_dataset(df_val_train, "val-training")
+        self.test_ds = self.torch_dataset(self.df_test, "test")
+
+    def get_collate_fn(self):
+        return None
+    
+    def get_batch_size(self, key:str) -> int:
+        return int(self.hypers.get(f'{key}_batch_size', self.hypers.get('batch_size', -1)))
+
+    def train_dataloader(self):
+        batch_size = self.get_batch_size('training')
+        return torch_data.DataLoader(self.train_ds, batch_size=batch_size, collate_fn=self.get_collate_fn(), shuffle=True,
+                                     num_workers=0, pin_memory=True)
+    
+    def _build_val_test_loader(self, ds):
+        batch_size = self.get_batch_size('val_test')
+        sampler = None
+        drop_last = False
+        if self.hypers.get('use_distributed_sampler', False) and util.is_distr_env():
+            sampler = torch_data.DistributedSampler(ds, drop_last=True, shuffle=False)
+            drop_last = True
+        dl = torch_data.DataLoader(ds, batch_size=batch_size, collate_fn=self.get_collate_fn(), sampler=sampler,
+                                   num_workers=0, pin_memory=True, drop_last=drop_last)
+        return dl
+
+    def val_dataloader(self):
+        train_dl = self._build_val_test_loader(self.val_train_ds)
+        val_dl = self._build_val_test_loader(self.val_ds)
+        return [train_dl, val_dl]
+    
+    def test_dataloader(self):
+        return self._build_val_test_loader(self.test_ds)
+    
+class RP3GlobalEmbeddingsLDM(RP3LDM):
+    def __init__(self, hypers) -> None:
+        super().__init__(hypers)
+        self.embeddings = None
+
+    def load_df(self):
+        df = super().load_df()
+        embeddings_file = self.hypers.embeddings_file
+        if embeddings_file == 'onehot' or embeddings_file.startswith('random_'):
+            seqs = util.read_fasta(self.hypers.fasta_path)
+            aa_to_int = {aa: i for i, aa in enumerate('ACDEFGHIKLMNPQRSTVWY')}
+            if embeddings_file.startswith('random_'):
+                emb = torch.nn.Embedding(20, int(embeddings_file[7:])).to('cpu').requires_grad_(False)
+            self.embeddings = dict()
+            for row in df.select('id', self.hypers.fasta_key).iter_rows():
+                seq = seqs[row[1]]
+                seq_tz = torch.tensor([aa_to_int[aa] for aa in seq], dtype=torch.int64)
+                if embeddings_file == 'onehot':
+                    seq_enc = torch.nn.functional.one_hot(seq_tz, num_classes=len(aa_to_int)).to(dtype=torch.float32)
+                else:
+                    seq_enc = emb(seq_tz)
+                self.embeddings[row[0]] = seq_enc.mean(0)
+        else:
+            embeddings_file = util.resolve(self.hypers.embeddings_file)
+            self.embeddings = load_global_embeddings_file(embeddings_file)
+        return df
+
+    def torch_dataset(self, df:pl.DataFrame, prefix:str) -> torch_data.Dataset:
+        return RP3GlobalEmbeddingsDataSet(df, prefix, self.embeddings)
+    
+class RP3SequenceEmbeddingsLDM(RP3LDM):
+    def __init__(self, hypers) -> None:
+        super().__init__(hypers)
+        self.embeddings_file = None
+
+    def load_df(self):
+        df = super().load_df()
+        self.embeddings_file = zipfile.ZipFile(util.resolve(self.hypers.embeddings_file), 'r')
+        return df
+    
+    @staticmethod
+    def collate(batch):
+        embeddings = [b.pop('embeddings') for b in batch]
+        ret = torch_data.default_collate(batch)
+        emb_len = torch.tensor([e.shape[0] for e in embeddings])
+        max_len = emb_len.max()
+        emb_padded = torch.stack([torch.nn.functional.pad(e, (0,0,0, max_len - e.shape[0]), value=0) for e in embeddings])
+        attn_mask = torch.zeros((emb_len.shape[0], max_len), dtype=torch.int32)
+        for i, l in enumerate(emb_len):
+            attn_mask[i, :l] = 1
+        ret['embeddings'] = emb_padded
+        ret['attention_mask'] = attn_mask
+        return ret
+    
+    def get_collate_fn(self):
+        return RP3SequenceEmbeddingsLDM.collate
+    
+    def torch_dataset(self, df:pl.DataFrame, prefix:str) -> torch_data.Dataset:
+        return RP3SequenceEmbeddingsDataSet(df, prefix, self.embeddings_file)
+    
+
+class RP3SequenceLDM(RP3LDM):
+
+    def __init__(self, hypers) -> None:
+        super().__init__(hypers)
+
+    def load_df(self):
+        df = super().load_df()
+        log.info(f"Reading sequences from {self.hypers.fasta_path}")
+        fasta_map = util.read_fasta(self.hypers.fasta_path)
+        fasta_id_col = self.hypers.fasta_id_col
+        df = df.with_columns(seq=pl.col(fasta_id_col).replace_strict(fasta_map))
+        return df
+    
+    def torch_dataset(self, df:pl.DataFrame, prefix:str) -> torch_data.Dataset:
+        return RP3SequenceDataSet(df, prefix, self.rng, self.hypers.get('max_seq_len', 0))
+
+    @staticmethod
+    def collate(tokenizer:model.RP3Net, batch):
+        seqs = [b.pop('seq') for b in batch]
+        ret = torch_data.default_collate(batch)
+        ret['seq'] = tokenizer.tokenize_sequences(seqs)
+        return ret
+    
+    def get_collate_fn(self):
+        lm: RP3LM = self.trainer.lightning_module
+        return functools.partial(RP3SequenceLDM.collate, lm.model)
+    

RP3Net/training/data_emlc.py

@@ -0,0 +1,94 @@
+import logging
+import numpy as np
+import polars as pl
+import torch.utils.data as torch_data
+
+from . import data
+from .. import util
+
+log = util.get_logger(__name__)
+
+class EmlcBatchSampler(torch_data.Sampler):
+    def __init__(self, *, df:pl.DataFrame, rng:np.random.Generator, clean_sources:list[int], noisy_sources:list[int],
+                  batch_size_clean:int, emlc_k:int=1, world_size:int=1, global_rank:int=0):
+        self.df = df.with_row_index('_row_idx')
+        self.rng = rng
+        self.emlc_k = emlc_k
+        self.clean_sources = clean_sources
+        self.noisy_sources = noisy_sources
+        self.batch_size_clean = batch_size_clean
+        self.batch_size_noisy = batch_size_clean * emlc_k
+        self.batch_count = min(df.select(pl.col('source').is_in(clean_sources)).sum()[0,0] // (self.batch_size_clean * world_size),
+                               df.select(pl.col.source.is_in(noisy_sources)).sum()[0,0] // (self.batch_size_noisy * world_size))
+        self.global_rank = global_rank
+        self.world_size = world_size
+        
+    def __iter__(self):
+        ix_clean = self.df.filter(pl.col.source.is_in(self.clean_sources)).select('_row_idx').to_numpy().flatten()
+        self.rng.shuffle(ix_clean)
+        ix_noisy = self.df.filter(pl.col.source.is_in(self.noisy_sources)).select('_row_idx').to_numpy().flatten()
+        self.rng.shuffle(ix_noisy)
+        if log.isEnabledFor(logging.DEBUG):
+            log.debug(f"Clean index:\n{ix_clean[:100]}")
+            log.debug(f"Noisy index:\n{ix_noisy[:100]}")
+        ix_clean = ix_clean[self.global_rank::self.world_size]
+        ix_noisy = ix_noisy[self.global_rank::self.world_size]
+        log.debug(f"{self.batch_count} batches per worker {self.global_rank}/{self.world_size}")
+        for i in range(self.batch_count):
+            out = np.concatenate([ix_clean[i*self.batch_size_clean:(i+1)*self.batch_size_clean],
+                                  ix_noisy[i*self.batch_size_noisy:(i+1)*self.batch_size_noisy]])
+            log.debug(f"Batch {i}/{self.batch_count}: {out}")
+            yield out
+
+    def __len__(self):
+        return self.batch_count
+    
+
+class EmlcLDM(data.RP3SequenceLDM):
+    def __init__(self, hypers) -> None:
+        super().__init__(hypers)
+        log.info("EmlcLDM init")
+        clean_sources = self.hypers.clean_sources
+        self.emlc_k = int(self.hypers.emlc_k)
+        self.clean_sources = [self.sources_map[s] for s in clean_sources]
+
+    def create_torch_datasets(self):
+        self.train_ds = self.torch_dataset(self.df_train, 'train')
+        self.val_ds = self.torch_dataset(self.df_val, "val")
+        df_val_train = (self.df_train
+            .filter(pl.col.source.is_in(self.clean_sources))
+            .sample(len(self.val_ds), with_replacement=False, seed=self.hypers.test_val_seed)
+        )
+        self.val_train_ds = self.torch_dataset(df_val_train, "val-training")
+        self.test_ds = self.torch_dataset(self.df_test, "test")
+
+    def setup(self, stage: str) -> None:
+        log.debug("EmlcLDM setup")
+        if self.trainer is not None:
+            assert self.trainer.model.sources == self.hypers.sources
+            assert self.trainer.model.emlc_k == self.emlc_k
+        df = self.load_df()
+        self.df_train = df.filter(
+            pl.col('ds_type').is_not_null() &
+            pl.col('ds_type').is_in(['TEST', self.validation_slice]).not_()
+        )
+        assert self.df_train.shape[0] > 0, f"No training data for slice {self.validation_slice}"
+        self.df_val = df.filter(pl.col('ds_type') == self.validation_slice, pl.col.source.is_in(self.clean_sources))
+        assert self.df_val.shape[0] > 0, f"No validation data for slice {self.validation_slice}"
+        self.df_test = df.filter(pl.col('ds_type') == 'TEST', pl.col.source.is_in(self.clean_sources))
+        assert self.df_test.shape[0] > 0, f"No test data for slice {self.validation_slice}"
+        self.create_torch_datasets()
+
+        
+    def train_dataloader(self):
+        log.debug("EmlcCVLDM train_dataloader")
+        batch_size = self.get_batch_size('training')
+        noisy_sources = [self.sources_map[s] for s in self.sources_map if self.sources_map[s] not in self.clean_sources]
+        world_size = self.trainer.world_size
+        global_rank = self.trainer.global_rank
+        batch_sampler = EmlcBatchSampler(df=self.df_train, rng=self.rng, 
+                                         clean_sources=self.clean_sources, noisy_sources=noisy_sources, 
+                                         batch_size_clean=batch_size, emlc_k=self.emlc_k,
+                                         world_size=world_size, global_rank=global_rank)
+        return torch_data.DataLoader(self.train_ds, batch_sampler=batch_sampler, num_workers=0, pin_memory=True, collate_fn=self.get_collate_fn())
+    

RP3Net/training/lm.py

@@ -0,0 +1,123 @@
+import lightning as L
+import lightning.pytorch.utilities as L_util
+
+import torch.nn as nn
+import torch
+import ml_collections as mlc
+import pandas as pd
+import os
+
+from . import metrics
+from .. import util
+from .. import model
+
+
+log = util.get_logger(__name__)
+
+class RP3LM(L.LightningModule):
+    def __init__(self, hypers) -> None:
+        super().__init__()
+        log.debug("Lightning module init")
+        self._hypers_prefix = 'model'
+        self.save_hyperparameters({'model': hypers})
+        self.hypers = mlc.ConfigDict(self.hparams.model)
+        self.sources = self.hypers.sources
+        self.sources_map = {s:i for i, s in enumerate(self.sources)}
+        log.info(f"Sources: {self.sources}")
+        self.metrics = metrics.ClassificationMetricContainer.create_classification_metrics(self.sources, 2)
+        self.loss = nn.CrossEntropyLoss()
+        log.info(f"Loss: {self.loss}")
+        self.model: model.RP3Net = model.load_model(self.hypers.model)
+        log.info(f"Model: {self.model}")
+
+    def setup(self, stage):
+        if stage == 'fit':
+            assert self.model.mode in model.Mode_Training, "Model must be in training mode"
+
+    def force_train_on_fit_start(self):
+        """
+        Need this, because loading a pre-trained HF model calls .eval() under the hood,
+        and PL preserves the state of training flags on modules when switching back from eval to train.
+        """
+        self.model.train()
+
+    def on_fit_start(self) -> None:
+        self.force_train_on_fit_start()
+    
+    def forward(self, batch):
+        return self.model(batch)
+        
+    def predict_step(self, batch, batch_idx: int, dataloader_idx: int = 0):
+        logits = self.model(batch)
+        return torch.argmax(logits, dim=1)
+    
+    def training_step(self, batch, batch_idx):
+        log.debug(f"Training batch ids: {batch['idx']}")        
+        logits = self(batch)
+        loss = self.loss(logits, batch['yield_binary'])
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch['yield_binary'].shape[0])
+        return loss 
+
+    def validation_step(self, batch, batch_idx, dataloader_idx):
+        ids = batch['idx']
+        log.debug(f"Validation batch ids for dataloader index {dataloader_idx}: {ids}")
+        logits = self(batch)
+        if dataloader_idx == 0:
+            self.metrics.update_train(logits, batch)
+        elif dataloader_idx == 1:
+            self.metrics.update_val(logits, batch)
+        else:
+            raise RuntimeError(f"Unknown dataloader index: {dataloader_idx}")
+
+    def test_step(self, batch, batch_idx):
+        log.debug(f"Test batch index: {batch['idx']}")
+        logits = self(batch)
+        self.metrics.update_test(logits, batch)
+
+    @L_util.rank_zero_only
+    def write_results_df(self, filename:os.PathLike, ids:torch.Tensor, logits:torch.Tensor):
+        proba = torch.softmax(logits, axis=1).cpu().numpy()
+        y_hat = proba.argmax(axis=1)
+        df = pd.DataFrame({'id': ids.to(dtype=torch.int32, device='cpu').numpy(), 'y_hat': y_hat})
+        df_logits = pd.DataFrame(logits.cpu().numpy(), columns=[f'logit_{i}' for i in range(logits.shape[1])])
+        df_proba = pd.DataFrame(proba, columns=[f'prob_{i}' for i in range(proba.shape[1])])
+        df = pd.concat([df, df_logits, df_proba], axis=1)
+        df.to_csv(filename, index=False)
+
+    def on_validation_epoch_end(self) -> None:
+        log.info(f"Validation epoch {self.current_epoch} end.")
+        train_log_dict = self.metrics.compute_train_dict()
+        if not self.trainer.sanity_checking:
+            self.log_dict(train_log_dict, on_epoch=True, add_dataloader_idx=False, sync_dist=True)
+        
+        val_log_dict = self.metrics.compute_val_dict()
+        if not self.trainer.sanity_checking:
+            self.log_dict(val_log_dict, on_epoch=True, add_dataloader_idx=False, sync_dist=True)
+        
+        train_df_file = util.resolve(self.trainer.default_root_dir) / f"train_df_{self.current_epoch}.csv.gz"
+        train_ids, train_logits = self.metrics.train_curve()
+        if isinstance(train_logits, torch.Tensor) and train_logits.shape[0] > 0 and not self.trainer.sanity_checking:
+            log.info(f"Writing training results for epoch {self.current_epoch} to {train_df_file}")
+            self.write_results_df(train_df_file, train_ids, train_logits)
+        
+        val_df_file = util.resolve(self.trainer.default_root_dir) / f"val_df_{self.current_epoch}.csv.gz"
+        val_ids, val_logits = self.metrics.val_curve()
+        if isinstance(val_logits, torch.Tensor) and val_logits.shape[0] > 0 and not self.trainer.sanity_checking:
+            log.info(f"Writing validation results for epoch {self.current_epoch} to {val_df_file}")
+            self.write_results_df(val_df_file, val_ids, val_logits)
+        self.metrics.reset()
+        
+
+    def on_test_epoch_end(self) -> None:
+        test_log_dict = self.metrics.compute_test_dict()
+        self.log_dict(test_log_dict, on_epoch=True, add_dataloader_idx=False)
+        test_df_file = util.resolve(self.trainer.default_root_dir) / f"test_df.csv.gz"
+        test_ids, test_logits = self.metrics.test_curve()
+        if isinstance(test_logits, torch.Tensor) and test_logits.shape[0] > 0:
+            log.info(f"Writing test results to {test_df_file}")
+            self.write_results_df(test_df_file, test_ids, test_logits)
+        self.metrics.reset()
+
+
+
+