def fetch()

in summarize_from_feedback/query_response_model.py [0:0]

• 16 lines of code
• 6 McCabe index (conditional complexity)

        def fetch(chkpt_prefix: str, module_name: str = ""):
            sharded_pieces = [
                fetch_single_piece(f"{chkpt_prefix}_shard_{shard_idx:03d}.pkl")
                for shard_idx in load_shard_idxs
            ]

            model_piece = {}
            for k in sharded_pieces[0].keys():
                parameter_name = ".".join([module_name, k]) if module_name else k

                sharding_dim = dim_to_shard(parameter_name)

                if sharding_dim is None:
                    val = sharded_pieces[0][k]
                else:
                    val = torch.cat([piece[k] for piece in sharded_pieces], dim=sharding_dim)
                fix_factor = get_shard_fix_factor(parameter_name, model_H, old_model_H)
                model_piece[k] = (val.float() * fix_factor).to(val.dtype)

            return model_piece