def load()

in summarize_from_feedback/query_response_model.py [0:0]

• 35 lines of code
• 11 McCabe index (conditional complexity)

    def load(self, load_path, run_params=None, init_heads=(), map_heads={}, use_cache=False):
        """
        Rebuilds everything, but keeps API semantics: model has same layout, and is on the same device, and all heads are the same (although some may be random init)
        """
        if use_cache and blobs.is_blob_url(load_path):
            load_path = blobs.download_directory_cached(load_path)

        with bf.BlobFile(os.path.join(load_path, "info.json")) as f:
            info = json.load(f)
        self.model_hparams = Hyperparams(info["model_hparams"])
        if run_params is not None:
            extra_model_H = {k: v for k, v in run_params.to_json().items() if v is not None}
            self.model_hparams.update(**extra_model_H)
        self.encoder = summarize_from_feedback.encoder
        model = build_with_random_weights(
            layout=self.layout,
            n_vocab=self.encoder.n_vocab,
            device=self.device,
            model_H=self.model_hparams,
        )

        self.model = self._update_model_with_head_info(model)

        init_heads = set(init_heads or ())
        # Load heads from where map_heads says, or the normal head name by default
        load_heads_map = {
            head: map_heads.get(head, head) for head in self.heads if head not in init_heads
        }
        load_exported_model(
            self.layout,
            self.model,
            self.model_hparams,
            load_path,
            load_heads_map=load_heads_map,
            use_cache=use_cache,
        )
        params_to_init = []
        for head in init_heads:
            params_to_init.append(self.model.scalar_heads[head].weight)
            params_to_init.append(self.model.scalar_heads[head].bias)

        self._sync_params(params_to_init, heads_to_init=init_heads)

        self.barrier("load_finished")