def fetch()

in summarize_from_feedback/query_response_model.py [0:0]

• 20 lines of code
• 4 McCabe index (conditional complexity)

        def fetch(chkpt_prefix: str, module_name: str = ""):
            unsharded_piece = fetch_single_piece(
                f"{chkpt_prefix}_shard_{shard_idx_to_load:03d}.pkl"
            )

            model_piece = {}
            for k in unsharded_piece.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 = unsharded_piece[k]
                else:
                    split_size = exact_div(
                        unsharded_piece[k].size()[sharding_dim], n_ranks_per_chkpt_shard
                    )
                    val = torch.split(
                        unsharded_piece[k], [split_size] * n_ranks_per_chkpt_shard, dim=sharding_dim
                    )[shard_slice_idx]
                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