steerable_retrieval/models/base.py

from torch import nn

class BaseModule(nn.Module):
    def __init__(self, ckpt_path = None, freeze = False):
        super().__init__()
        self.ckpt_path = ckpt_path
        self.freeze = freeze


         
    @classmethod
    def from_config(cls, config, device=None):
        config.update(device=device)
        return cls(**config)
    
    @classmethod
    def from_yaml(cls, yaml_path, device=None):
        
        if 's3://' in yaml_path:
            import s3fs
            fs = s3fs.S3FileSystem()
            with fs.open(yaml_path, "r") as file:
                config = yaml.safe_load(file)
        else:
            with open(yaml_path, "r") as file:
                config = yaml.safe_load(file)
        
        config = config.get('model', config)
        config = config.get('init_args', config)
        
        return cls.from_config(config, device=device)
    
    @classmethod
    def from_pretrained(cls, yaml_or_config, ckpt_path, device=None):
        if isinstance(yaml_or_config, str):
            model = cls.from_yaml(yaml_or_config, device=device)
        else:
            model = cls.from_config(yaml_or_config, device=device)
            
        if 's3://' in ckpt_path:
            from s3torchconnector import S3Checkpoint
            checkpoint= S3Checkpoint(region='us-east-1')
            with checkpoint.reader(ckpt_path) as f:
                ckpt = torch.load(f, map_location=device)
                model.load_state_dict(ckpt['state_dict'], strict=True)
                print(f"Model loaded from {ckpt_path}")
        else:
            ckpt = torch.load(ckpt_path, map_location=device)
            model.load_state_dict(ckpt['state_dict'], strict=True)
            print(f"Model loaded from {ckpt_path}")
        
        return model
    

    def configure_optimizers(self):
        from steerable_retrieval.utils.instantiators import instantiate
        
        if not hasattr(self, 'optimizer'):
            self.optimizer = None
        if self.optimizer is None:
            optimizer = optim.Adam(
                self.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-8)
        else:
            # If optimizer is a config dict or partial, instantiate it with model parameters
            if isinstance(self.optimizer, dict):
                # Handle _partial_ configs - Hydra creates a functools.partial
                optimizer_cfg = dict(self.optimizer)
                
                # If using _target_ style (with or without _partial_)
                if '_target_' in optimizer_cfg:
                    # Remove _partial_ flag if present (it was just to prevent instantiation)
                    optimizer_cfg.pop('_partial_', None)
                    optimizer_cfg.pop('_convert_', None)
                    optimizer_cfg['params'] = self.parameters()
                    optimizer = instantiate(optimizer_cfg)
                # If using class_path style
                elif 'class_path' in optimizer_cfg:
                    import importlib
                    module_path, class_name = optimizer_cfg['class_path'].rsplit('.', 1)
                    module = importlib.import_module(module_path)
                    optimizer_class = getattr(module, class_name)
                    init_args = optimizer_cfg.get('init_args', {})
                    init_args['params'] = self.parameters()
                    optimizer = optimizer_class(**init_args)
                else:
                    # Fallback: assume it's a direct config
                    optimizer_cfg['params'] = self.parameters()
                    optimizer = instantiate(optimizer_cfg)
            elif hasattr(self.optimizer, 'func') and hasattr(self.optimizer, 'keywords'):
                # It's a functools.partial (from _partial_=true) - call it with params
                optimizer = self.optimizer(params=self.parameters())
            elif callable(self.optimizer):
                # If it's a callable (old style), call it with parameters
                optimizer = self.optimizer(self.parameters())
            else:
                # Fallback to default
                optimizer = optim.Adam(
                    self.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-8)
            
        if hasattr(self, 'scheduler') and self.scheduler is not None:
            # copy of the scheduler applied to the optimizer
            ## retrocompatibilty with old schedulers
            if isinstance(self.scheduler, dict) and 'class_name' in self.scheduler.keys():
                scheduler_class = eval(self.scheduler['class_name'])
                scheduler_kwargs = self.scheduler.get('init_args', {})
                scheduler = scheduler_class(optimizer, **scheduler_kwargs)
                self.scheduler = scheduler  # Store instantiated scheduler
                # Return with proper Lightning configuration
                return {
                    'optimizer': optimizer,
                    'lr_scheduler': {
                        'scheduler': scheduler,
                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
                        'frequency': 1,      # Step every optimizer step
                    }
                }
            elif isinstance(self.scheduler, dict):
                # Handle configs that were prevented from instantiation
                scheduler_cfg = dict(self.scheduler)
                
                # If using class_path style (not _target_)
                if 'class_path' in scheduler_cfg:
                    import importlib
                    module_path, class_name = scheduler_cfg['class_path'].rsplit('.', 1)
                    module = importlib.import_module(module_path)
                    scheduler_class = getattr(module, class_name)
                    init_args = scheduler_cfg.get('init_args', {})
                    init_args['optimizer'] = optimizer
                    scheduler = scheduler_class(**init_args)
                # If using _target_ style
                elif '_target_' in scheduler_cfg:
                    # Remove _partial_ flag and add optimizer
                    scheduler_cfg.pop('_partial_', None)
                    scheduler_cfg.pop('_convert_', None)
                    scheduler_cfg['optimizer'] = optimizer
                    scheduler = instantiate(scheduler_cfg)
                else:
                    # Fallback: assume it's a direct config
                    scheduler_cfg['optimizer'] = optimizer
                    scheduler = instantiate(scheduler_cfg)
                self.scheduler = scheduler  # Store instantiated scheduler
                # Return with proper Lightning configuration
                return {
                    'optimizer': optimizer,
                    'lr_scheduler': {
                        'scheduler': scheduler,
                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
                        'frequency': 1,      # Step every optimizer step
                    }
                }
            elif hasattr(self.scheduler, 'func') and hasattr(self.scheduler, 'keywords'):
                # It's a functools.partial (from _partial_=true) - call it with optimizer
                scheduler = self.scheduler(optimizer=optimizer)
                self.scheduler = scheduler  # Store instantiated scheduler
                # Return with proper Lightning configuration
                return {
                    'optimizer': optimizer,
                    'lr_scheduler': {
                        'scheduler': scheduler,
                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
                        'frequency': 1,      # Step every optimizer step
                    }
                }
            else:
                # If scheduler is already instantiated, just return it
                # Return with proper Lightning configuration
                return {
                    'optimizer': optimizer,
                    'lr_scheduler': {
                        'scheduler': self.scheduler,
                        'interval': 'step',  # Step after optimizer steps (respects accumulate_grad_batches)
                        'frequency': 1,      # Step every optimizer step
                    }
                }
        
        return optimizer

    def load_ckpt(self, ckpt_path, device = None, prefix = ''):

        if device is None:
            device = next(self.parameters()).device
            
        if 's3://' in ckpt_path:
            from s3torchconnector import S3Checkpoint
            checkpoint= S3Checkpoint(region='us-east-1')
            with checkpoint.reader(ckpt_path) as f:
                state_dict = torch.load(f, map_location=device)['state_dict']
                print(f"Model loaded from {ckpt_path}")
        else:
            state_dict = torch.load(ckpt_path, map_location=device)['state_dict']
            print(f"Model loaded from {ckpt_path}")
        
        print(state_dict)
        
        try:
            self.load_state_dict(state_dict)
            print("Loaded full state dict")
        except:
            print("Could not load state dict, trying to load only ['encoder'] keys")
            
            try:
                from collections import OrderedDict
                new_state_dict = OrderedDict()
                for k in list(state_dict.keys()):
                    if prefix in k:
                        new_key = k.replace('encoder.','')
                        new_state_dict[new_key] = state_dict[k]
                        
                self.load_state_dict(new_state_dict)
                print(f"Loaded only {prefix} keys")
                
            except Exception as e:
                print(f"Could not load state dict, error: {e}")
            
        

    def freeze(self):
        for param in self.parameters():
            param.requires_grad = False