def __init__()

in summarize_from_feedback/query_response_model.py [0:0]

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

    def __init__(self, model, scalar_heads, d_model, init_scales=1.0):
        super().__init__()
        self.model = model
        self.scalar_head_names = scalar_heads
        if not isinstance(init_scales, dict):
            init_scales = {head_name: init_scales for head_name in scalar_heads}

        self.scalar_heads = torch.nn.ModuleDict()
        for name in self.scalar_head_names:
            head = torch.nn.Linear(d_model, 1)
            init_std = init_scales.get(name, 1.0) / np.sqrt(d_model + 1)
            torch.nn.init.normal_(head.weight, std=init_std)
            torch.nn.init.zeros_(head.bias)
            self.scalar_heads[name] = head

        for attr in [
            "include_input_embeddings",
            "embedding",
            "include_pos_embeddings",
            "position_embedding",
            "include_final_layer_norm",
            "include_output_unembeddings",
            "ln_f",
            "unembedding_weights",
            "torso",
            "mp_comm",
            "n_ctx",
        ]:
            if hasattr(self.model, attr):
                setattr(self, attr, getattr(self.model, attr))