punctuation-knowledge-neurons.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install einops\n# main knowledge neurons class\nimport torch\nimport torch.nn.functional as F\nimport torch.nn as nn\nimport einops\nfrom tqdm import tqdm\nimport numpy as np\nimport collections\nfrom typing import List, Optional, Tuple, Callable\nimport string\nimport math\nfrom functools import partial\nfrom transformers import PreTrainedTokenizerBase\nimport random\nfrom functools import lru_cache\nfrom functools import reduce\nimport operator\n\nimport json\nfrom pathlib import Path \nimport os\nfrom glob import glob\nimport seaborn as sns\nimport pandas as pd \n\nfrom transformers import (\n    BertTokenizer,\n    BertLMHeadModel,\n    GPT2Tokenizer,\n    GPT2LMHeadModel,\n)\n\nimport urllib.request\n","metadata":{"execution":{"iopub.status.busy":"2022-12-13T01:27:04.709375Z","iopub.execute_input":"2022-12-13T01:27:04.709841Z","iopub.status.idle":"2022-12-13T01:27:24.058313Z","shell.execute_reply.started":"2022-12-13T01:27:04.709806Z","shell.execute_reply":"2022-12-13T01:27:24.056254Z"},"trusted":true},"execution_count":2,"outputs":[{"name":"stdout","text":"Collecting einops\n  Downloading einops-0.6.0-py3-none-any.whl (41 kB)\n\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.6/41.6 kB\u001b[0m \u001b[31m176.9 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n\u001b[?25hInstalling collected packages: einops\nSuccessfully installed einops-0.6.0\n\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m","output_type":"stream"}]},{"cell_type":"code","source":"\nBERT_MODELS = [\"bert-base-uncased\"]\nGPT2_MODELS = [\"gpt2\"]\nALL_MODELS = BERT_MODELS + GPT2_MODELS\n\n\ndef initialize_model_and_tokenizer(model_name: str):\n    if model_name in BERT_MODELS:\n        tokenizer = BertTokenizer.from_pretrained(model_name)\n        model= BertLMHeadModel.from_pretrained(model_name)\n    elif model_name in GPT2_MODELS:\n        tokenizer = GPT2Tokenizer.from_pretrained(model_name)\n        model = GPT2LMHeadModel.from_pretrained(model_name)\n    else:\n        raise ValueError(\"Model {model_name} not supported\")\n\n    model.eval()\n\n    return model, tokenizer\n\n\ndef model_type(model_name: str):\n    if model_name in BERT_MODELS:\n        return \"bert\"\n    elif model_name in GPT2_MODELS:\n        return \"gpt2\"\n    else:\n        raise ValueError(\"Model {model_name} not supported\")\n","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:03:08.437899Z","iopub.execute_input":"2022-11-28T16:03:08.439383Z","iopub.status.idle":"2022-11-28T16:03:08.447211Z","shell.execute_reply.started":"2022-11-28T16:03:08.439341Z","shell.execute_reply":"2022-11-28T16:03:08.446145Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"code","source":"class KnowledgeNeurons:\n    def __init__(\n        self,\n        model: nn.Module,\n        tokenizer: PreTrainedTokenizerBase,\n        model_type: str = \"bert\",\n        device: str = None,\n    ):\n        self.model = model\n        self.model_type = model_type\n        self.device = device or torch.device(\n            \"cuda\" if torch.cuda.is_available() else \"cpu\"\n        )\n        self.model.to(self.device)\n        self.tokenizer = tokenizer\n\n        self.baseline_activations = None\n\n        if self.model_type == \"bert\":\n            self.transformer_layers_attr = \"bert.encoder.layer\"\n            self.input_ff_attr = \"intermediate\"\n            self.output_ff_attr = \"output.dense.weight\"\n            self.word_embeddings_attr = \"bert.embeddings.word_embeddings.weight\"\n            self.unk_token = getattr(self.tokenizer, \"unk_token_id\", None)\n        elif \"gpt2\" in model_type:\n            self.transformer_layers_attr = \"transformer.h\"\n            self.input_ff_attr = \"mlp.c_fc\"\n            self.output_ff_attr = \"mlp.c_proj.weight\"\n            self.word_embeddings_attr = \"transformer.wpe\"\n        else:\n            raise NotImplementedError\n\n    def _get_output_ff_layer(self, layer_idx):\n        return get_ff_layer(\n            self.model,\n            layer_idx,\n            transformer_layers_attr=self.transformer_layers_attr,\n            ff_attrs=self.output_ff_attr,\n        )\n\n    def _get_input_ff_layer(self, layer_idx):\n        return get_ff_layer(\n            self.model,\n            layer_idx,\n            transformer_layers_attr=self.transformer_layers_attr,\n            ff_attrs=self.input_ff_attr,\n        )\n\n    def _get_word_embeddings(self):\n        return get_attributes(self.model, self.word_embeddings_attr)\n\n    def _get_transformer_layers(self):\n        return get_attributes(self.model, self.transformer_layers_attr)\n\n    def _prepare_inputs(self, prompt, target=None, encoded_input=None):\n        if encoded_input is None:\n            encoded_input = self.tokenizer(prompt, return_tensors=\"pt\").to(self.device)\n        if self.model_type == \"bert\":\n            mask_idx = torch.where(\n                encoded_input[\"input_ids\"][0] == self.tokenizer.mask_token_id\n            )[0].item()\n        else:\n            # with autoregressive models we always want to target the last token\n            mask_idx = -1\n        if target is not None:\n            if \"gpt2\" in self.model_type:\n                target = self.tokenizer.encode(target)\n            else:\n                target = self.tokenizer.convert_tokens_to_ids(target)\n        return encoded_input, mask_idx, target\n\n    def _generate(self, prompt, ground_truth):\n        encoded_input, mask_idx, target_label = self._prepare_inputs(\n            prompt, ground_truth\n        )\n        # for autoregressive models, we might want to generate > 1 token\n        if self.model_type == \"gpt2\":\n            n_sampling_steps = len(target_label)\n        else:\n            n_sampling_steps = 1  # TODO: we might want to use multiple mask tokens even with bert models\n\n        all_gt_probs = []\n        all_argmax_probs = []\n        argmax_tokens = []\n        argmax_completion_str = \"\"\n\n        for i in range(n_sampling_steps):\n            if i > 0:\n                # retokenize new inputs\n                encoded_input, mask_idx, target_label = self._prepare_inputs(\n                    prompt, ground_truth\n                )\n            outputs = self.model(**encoded_input)\n            probs = F.softmax(outputs.logits[:, mask_idx, :], dim=-1)\n            if n_sampling_steps > 1:\n                target_idx = target_label[i]\n            else:\n                target_idx = target_label\n            gt_prob = probs[:, target_idx].item()\n            all_gt_probs.append(gt_prob)\n\n            # get info about argmax completion\n            argmax_prob, argmax_id = [i.item() for i in probs.max(dim=-1)]\n            argmax_tokens.append(argmax_id)\n            argmax_str = self.tokenizer.decode([argmax_id])\n            all_argmax_probs.append(argmax_prob)\n\n            prompt += argmax_str\n            argmax_completion_str += argmax_str\n\n        gt_prob = reduce(operator.mul, all_gt_probs, 1) if len(all_gt_probs) > 1 else all_gt_probs[0]\n        argmax_prob = (\n            reduce(operator.mul, all_argmax_probs, 1)\n            if len(all_argmax_probs) > 1\n            else all_argmax_probs[0]\n        )\n        return gt_prob, argmax_prob, argmax_completion_str, argmax_tokens\n\n    def n_layers(self):\n        return len(self._get_transformer_layers())\n\n    def intermediate_size(self):\n        if self.model_type == \"bert\":\n            return self.model.config.intermediate_size\n        else:\n            return self.model.config.hidden_size * 4\n\n    @staticmethod\n    def scaled_input(activations: torch.Tensor, steps: int = 20, device: str = \"cpu\"):\n        \"\"\"\n        Tiles activations along the batch dimension - gradually scaling them over\n        `steps` steps from 0 to their original value over the batch dimensions.\n        `activations`: torch.Tensor\n        original activations\n        `steps`: int\n        number of steps to take\n        \"\"\"\n        tiled_activations = einops.repeat(activations, \"b d -> (r b) d\", r=steps)\n        out = (\n            tiled_activations\n            * torch.linspace(start=0, end=1, steps=steps).to(device)[:, None]\n        )\n        return out\n\n    def get_baseline_with_activations(\n        self, encoded_input: dict, layer_idx: int, mask_idx: int\n    ):\n        \"\"\"\n        Gets the baseline outputs and activations for the unmodified model at a given index.\n        `encoded_input`: torch.Tensor\n            the inputs to the model from self.tokenizer.encode_plus()\n        `layer_idx`: int\n            which transformer layer to access\n        `mask_idx`: int\n            the position at which to get the activations (TODO: rename? with autoregressive models there's no mask, so)\n        \"\"\"\n\n        def get_activations(model, layer_idx, mask_idx):\n            \"\"\"\n            This hook function should assign the intermediate activations at a given layer / mask idx\n            to the 'self.baseline_activations' variable\n            \"\"\"\n\n            def hook_fn(acts):\n                self.baseline_activations = acts[:, mask_idx, :]\n\n            return register_hook(\n                model,\n                layer_idx=layer_idx,\n                f=hook_fn,\n                transformer_layers_attr=self.transformer_layers_attr,\n                ff_attrs=self.input_ff_attr,\n            )\n\n        handle = get_activations(self.model, layer_idx=layer_idx, mask_idx=mask_idx)\n        baseline_outputs = self.model(**encoded_input)\n        handle.remove()\n        baseline_activations = self.baseline_activations\n        self.baseline_activations = None\n        return baseline_outputs, baseline_activations\n\n    def get_scores(\n        self,\n        prompt: str,\n        ground_truth: str,\n        batch_size: int = 10,\n        steps: int = 20,\n        attribution_method: str = \"integrated_grads\",\n        pbar: bool = True,\n    ):\n        \"\"\"\n        Gets the attribution scores for a given prompt and ground truth.\n        `prompt`: str\n            the prompt to get the attribution scores for\n        `ground_truth`: str\n            the ground truth / expected output\n        `batch_size`: int\n            batch size\n        `steps`: int\n            total number of steps (per token) for the integrated gradient calculations\n        `attribution_method`: str\n            the method to use for getting the scores. Choose from 'integrated_grads' or 'max_activations'.\n        \"\"\"\n\n        scores = []\n        encoded_input = self.tokenizer(prompt, return_tensors=\"pt\").to(self.device)\n        for layer_idx in tqdm(\n            range(self.n_layers()),\n            desc=\"Getting attribution scores for each layer...\",\n            disable=not pbar,\n        ):\n            layer_scores = self.get_scores_for_layer(\n                prompt,\n                ground_truth,\n                encoded_input=encoded_input,\n                layer_idx=layer_idx,\n                batch_size=batch_size,\n                steps=steps,\n                attribution_method=attribution_method,\n            )\n            scores.append(layer_scores)\n        return torch.stack(scores)\n\n    def get_coarse_neurons(\n        self,\n        prompt: str,\n        ground_truth: str,\n        batch_size: int = 10,\n        steps: int = 20,\n        threshold: float = None,\n        adaptive_threshold: float = None,\n        percentile: float = None,\n        attribution_method: str = \"integrated_grads\",\n        pbar: bool = True,\n    ) -> List[List[int]]:\n        \"\"\"\n        Finds the 'coarse' neurons for a given prompt and ground truth.\n        The coarse neurons are the neurons that are most activated by a single prompt.\n        We refine these by using multiple prompts that express the same 'fact'/relation in different ways.\n        `prompt`: str\n            the prompt to get the coarse neurons for\n        `ground_truth`: str\n            the ground truth / expected output\n        `batch_size`: int\n            batch size\n        `steps`: int\n            total number of steps (per token) for the integrated gradient calculations\n        `threshold`: float\n            `t` from the paper. If not None, then we only keep neurons with integrated grads above this threshold.\n        `adaptive_threshold`: float\n            Adaptively set `threshold` based on `maximum attribution score * adaptive_threshold` (in the paper, they set adaptive_threshold=0.3)\n        `percentile`: float\n            If not None, then we only keep neurons with integrated grads in this percentile of all integrated grads.\n        `attribution_method`: str\n            the method to use for getting the scores. Choose from 'integrated_grads' or 'max_activations'.\n        \"\"\"\n        attribution_scores = self.get_scores(\n            prompt,\n            ground_truth,\n            batch_size=batch_size,\n            steps=steps,\n            pbar=pbar,\n            attribution_method=attribution_method,\n        )\n        assert (\n            sum(e is not None for e in [threshold, adaptive_threshold, percentile]) == 1\n        ), f\"Provide one and only one of threshold / adaptive_threshold / percentile\"\n        if adaptive_threshold is not None:\n            threshold = attribution_scores.max().item() * adaptive_threshold\n        if threshold is not None:\n            return torch.nonzero(attribution_scores > threshold).cpu().tolist()\n        else:\n            s = attribution_scores.flatten().detach().cpu().numpy()\n            return (\n                torch.nonzero(attribution_scores > np.percentile(s, percentile))\n                .cpu()\n                .tolist()\n            )\n\n    def get_refined_neurons(\n        self,\n        prompts: List[str],\n        ground_truth: str,\n        negative_examples: Optional[List[str]] = None,\n        p: float = 0.5,\n        batch_size: int = 10,\n        steps: int = 20,\n        coarse_adaptive_threshold: Optional[float] = 0.3,\n        coarse_threshold: Optional[float] = None,\n        coarse_percentile: Optional[float] = None,\n        quiet=False,\n    ) -> List[List[int]]:\n        \"\"\"\n        Finds the 'refined' neurons for a given set of prompts and a ground truth / expected output.\n        The input should be n different prompts, each expressing the same fact in different ways.\n        For each prompt, we calculate the attribution scores of each intermediate neuron.\n        We then set an attribution score threshold, and we keep the neurons that are above this threshold.\n        Finally, considering the coarse neurons from all prompts, we set a sharing percentage threshold, p,\n        and retain only neurons shared by more than p% of prompts.\n        `prompts`: list of str\n            the prompts to get the refined neurons for\n        `ground_truth`: str\n            the ground truth / expected output\n        `negative_examples`: list of str\n            Optionally provide a list of negative examples. Any neuron that appears in these examples will be excluded from the final results.\n        `p`: float\n            the threshold for the sharing percentage\n        `batch_size`: int\n            batch size\n        `steps`: int\n            total number of steps (per token) for the integrated gradient calculations\n        `coarse_threshold`: float\n            threshold for the coarse neurons\n        `coarse_percentile`: float\n            percentile for the coarse neurons\n        \"\"\"\n        assert isinstance(\n            prompts, list\n        ), \"Must provide a list of different prompts to get refined neurons\"\n        assert 0.0 <= p <= 1.0, \"p should be a float between 0 and 1\"\n\n        n_prompts = len(prompts)\n        coarse_neurons = []\n        for prompt in tqdm(\n            prompts, desc=\"Getting coarse neurons for each prompt...\", disable=quiet\n        ):\n            coarse_neurons.append(\n                self.get_coarse_neurons(\n                    prompt,\n                    ground_truth,\n                    batch_size=batch_size,\n                    steps=steps,\n                    adaptive_threshold=coarse_adaptive_threshold,\n                    threshold=coarse_threshold,\n                    percentile=coarse_percentile,\n                    pbar=False,\n                )\n            )\n        if negative_examples is not None:\n            negative_neurons = []\n            for negative_example in tqdm(\n                negative_examples,\n                desc=\"Getting coarse neurons for negative examples\",\n                disable=quiet,\n            ):\n                negative_neurons.append(\n                    self.get_coarse_neurons(\n                        negative_example,\n                        ground_truth,\n                        batch_size=batch_size,\n                        steps=steps,\n                        adaptive_threshold=coarse_adaptive_threshold,\n                        threshold=coarse_threshold,\n                        percentile=coarse_percentile,\n                        pbar=False,\n                    )\n                )\n        if not quiet:\n            total_coarse_neurons = sum([len(i) for i in coarse_neurons])\n            print(f\"\\n{total_coarse_neurons} coarse neurons found - refining\")\n        t = n_prompts * p\n        refined_neurons = {0.6: [], 0.7: [], 0.8: [], 0.9: [], 1.0: []}\n        c = collections.Counter()\n        for neurons in coarse_neurons:\n            for n in neurons:\n                c[tuple(n)] += 1\n        for i in range(6, 11):\n            prop= n_prompts* i* 1.0/ 10\n            for neuron, count in c.items():\n                if count > prop:\n                    refined_neurons[i* 1.0/ 10].append(list(neuron))\n\n        # filter out neurons that are in the negative examples\n        if negative_examples is not None:\n            for neuron in negative_neurons:\n                if neuron in refined_neurons:\n                    refined_neurons.remove(neuron)\n\n        total_refined_neurons = len(refined_neurons)\n        if not quiet:\n            print(f\"{total_refined_neurons} neurons remaining after refining\")\n        return refined_neurons\n\n    def get_scores_for_layer(\n        self,\n        prompt: str,\n        ground_truth: str,\n        layer_idx: int,\n        batch_size: int = 10,\n        steps: int = 20,\n        encoded_input: Optional[int] = None,\n        attribution_method: str = \"integrated_grads\",\n    ):\n        \"\"\"\n        get the attribution scores for a given layer\n        `prompt`: str\n            the prompt to get the attribution scores for\n        `ground_truth`: str\n            the ground truth / expected output\n        `layer_idx`: int\n            the layer to get the scores for\n        `batch_size`: int\n            batch size\n        `steps`: int\n            total number of steps (per token) for the integrated gradient calculations\n        `encoded_input`: int\n            if not None, then use this encoded input instead of getting a new one\n        `attribution_method`: str\n            the method to use for getting the scores. Choose from 'integrated_grads' or 'max_activations'.\n        \"\"\"\n        assert steps % batch_size == 0\n        n_batches = steps // batch_size\n\n        # First we take the unmodified model and use a hook to return the baseline intermediate activations at our chosen target layer\n        encoded_input, mask_idx, target_label = self._prepare_inputs(\n            prompt, ground_truth, encoded_input\n        )\n\n        # for autoregressive models, we might want to generate > 1 token\n        if self.model_type == \"gpt2\":\n            n_sampling_steps = len(target_label)\n\n        else:\n            n_sampling_steps = 1  # TODO: we might want to use multiple mask tokens even with bert models\n\n        if attribution_method == \"integrated_grads\":\n            integrated_grads = []\n        \n            for i in range(n_sampling_steps):\n                if i > 0 and self.model_type == \"gpt2\":\n                    # retokenize new inputs\n                    encoded_input, mask_idx, target_label = self._prepare_inputs(\n                        prompt, ground_truth\n                    )\n                (\n                    baseline_outputs,\n                    baseline_activations,\n                ) = self.get_baseline_with_activations(\n                    encoded_input, layer_idx, mask_idx\n                )\n                if n_sampling_steps > 1:\n                    argmax_next_token = (\n                        baseline_outputs.logits[:, mask_idx, :].argmax(dim=-1).item()\n                    )\n                    next_token_str = self.tokenizer.decode(argmax_next_token)\n\n                # Now we want to gradually change the intermediate activations of our layer from 0 -> their original value\n                # and calculate the integrated gradient of the masked position at each step\n                # we do this by repeating the input across the batch dimension, multiplying the first batch by 0, the second by 0.1, etc., until we reach 1\n                scaled_weights = self.scaled_input(\n                    baseline_activations, steps=steps, device=self.device\n                )\n                scaled_weights.requires_grad_(True)\n\n                integrated_grads_this_step = []  # to store the integrated gradients\n\n                for batch_weights in scaled_weights.chunk(n_batches):\n                    # we want to replace the intermediate activations at some layer, at the mask position, with `batch_weights`\n                    # first tile the inputs to the correct batch size\n                    inputs = {\n                        \"input_ids\": einops.repeat(\n                            encoded_input[\"input_ids\"], \"b d -> (r b) d\", r=batch_size\n                        ),\n                        \"attention_mask\": einops.repeat(\n                            encoded_input[\"attention_mask\"],\n                            \"b d -> (r b) d\",\n                            r=batch_size,\n                        ),\n                    }\n                    if self.model_type == \"bert\":\n                        inputs[\"token_type_ids\"] = einops.repeat(\n                            encoded_input[\"token_type_ids\"],\n                            \"b d -> (r b) d\",\n                            r=batch_size,\n                        )\n\n                    # then patch the model to replace the activations with the scaled activations\n                    patch_ff_layer(\n                        self.model,\n                        layer_idx=layer_idx,\n                        mask_idx=mask_idx,\n                        replacement_activations=batch_weights,\n                        transformer_layers_attr=self.transformer_layers_attr,\n                        ff_attrs=self.input_ff_attr,\n                    )\n\n                    # then forward through the model to get the logits\n                    outputs = self.model(**inputs)\n\n                    # then calculate the gradients for each step w/r/t the inputs\n                    probs = F.softmax(outputs.logits[:, mask_idx, :], dim=-1)\n                    if n_sampling_steps > 1:\n                        target_idx = target_label[i]\n                    else:\n                        target_idx = target_label\n                    \n                    grad = torch.autograd.grad(\n                        torch.unbind(probs[:, target_idx]), batch_weights\n                    )[0]\n                    grad = grad.sum(dim=0)\n                    integrated_grads_this_step.append(grad)\n\n                    unpatch_ff_layer(\n                        self.model,\n                        layer_idx=layer_idx,\n                        transformer_layers_attr=self.transformer_layers_attr,\n                        ff_attrs=self.input_ff_attr,\n                    )\n\n                # then sum, and multiply by W-hat / m\n                integrated_grads_this_step = torch.stack(\n                    integrated_grads_this_step, dim=0\n                ).sum(dim=0)\n                integrated_grads_this_step *= baseline_activations.squeeze(0) / steps\n                integrated_grads.append(integrated_grads_this_step)\n\n                if n_sampling_steps > 1:\n                    prompt += next_token_str\n            integrated_grads = torch.stack(integrated_grads, dim=0).sum(dim=0) / len(\n                integrated_grads\n            )\n            return integrated_grads\n        elif attribution_method == \"max_activations\":\n            activations = []\n            for i in range(n_sampling_steps):\n                if i > 0 and self.model_type == \"gpt2\":\n                    # retokenize new inputs\n                    encoded_input, mask_idx, target_label = self._prepare_inputs(\n                        prompt, ground_truth\n                    )\n                (\n                    baseline_outputs,\n                    baseline_activations,\n                ) = self.get_baseline_with_activations(\n                    encoded_input, layer_idx, mask_idx\n                )\n                activations.append(baseline_activations)\n                if n_sampling_steps > 1:\n                    argmax_next_token = (\n                        baseline_outputs.logits[:, mask_idx, :].argmax(dim=-1).item()\n                    )\n                    next_token_str = self.tokenizer.decode(argmax_next_token)\n                    prompt += next_token_str\n            activations = torch.stack(activations, dim=0).sum(dim=0) / len(activations)\n            return activations.squeeze(0)\n        else:\n            raise NotImplementedError\n\n    def modify_activations(\n        self,\n        prompt: str,\n        ground_truth: str,\n        neurons: List[List[int]],\n        mode: str = \"suppress\",\n        undo_modification: bool = True,\n        quiet: bool = False,\n    ) -> Tuple[dict, Callable]:\n        results_dict = {}\n        _, mask_idx, _ = self._prepare_inputs(\n            prompt, ground_truth\n        )  # just need to get the mask index for later - probably a better way to do this\n        # get the baseline probabilities of the groundtruth being generated + the argmax / greedy completion before modifying the activations\n        (\n            gt_baseline_prob,\n            argmax_baseline_prob,\n            argmax_completion_str,\n            _,\n        ) = self._generate(prompt, ground_truth)\n        if not quiet:\n            print(\n                f\"\\nBefore modification - groundtruth probability: {gt_baseline_prob}\\nArgmax completion: `{argmax_completion_str}`\\nArgmax prob: {argmax_baseline_prob}\\n\"\n            )\n        results_dict[\"before\"] = {\n            \"gt_prob\": gt_baseline_prob,\n            \"argmax_completion\": argmax_completion_str,\n            \"argmax_prob\": argmax_baseline_prob,\n        }\n\n        # patch model to suppress neurons\n        # store all the layers we patch so we can unpatch them later\n        all_layers = set([n[0] for n in neurons])\n\n        patch_ff_layer(\n            self.model,\n            mask_idx,\n            mode=mode,\n            neurons=neurons,\n            transformer_layers_attr=self.transformer_layers_attr,\n            ff_attrs=self.input_ff_attr,\n        )\n\n        # get the probabilities of the groundtruth being generated + the argmax / greedy completion after modifying the activations\n        new_gt_prob, new_argmax_prob, new_argmax_completion_str, _ = self._generate(\n            prompt, ground_truth\n        )\n        if not quiet:\n            print(\n                f\"\\nAfter modification - groundtruth probability: {new_gt_prob}\\nArgmax completion: `{new_argmax_completion_str}`\\nArgmax prob: {new_argmax_prob}\\n\"\n            )\n        results_dict[\"after\"] = {\n            \"gt_prob\": new_gt_prob,\n            \"argmax_completion\": new_argmax_completion_str,\n            \"argmax_prob\": new_argmax_prob,\n        }\n\n        unpatch_fn = partial(\n            unpatch_ff_layers,\n            model=self.model,\n            layer_indices=all_layers,\n            transformer_layers_attr=self.transformer_layers_attr,\n            ff_attrs=self.input_ff_attr,\n        )\n\n        if undo_modification:\n            unpatch_fn()\n            unpatch_fn = lambda *args: args\n\n        return results_dict, unpatch_fn\n\n    def suppress_knowledge(\n        self,\n        prompt: str,\n        ground_truth: str,\n        neurons: List[List[int]],\n        undo_modification: bool = True,\n        quiet: bool = False,\n    ) -> Tuple[dict, Callable]:\n        \"\"\"\n        prompt the model with `prompt`, zeroing the activations at the positions specified by `neurons`,\n        and measure the resulting affect on the ground truth probability.\n        \"\"\"\n        return self.modify_activations(\n            prompt=prompt,\n            ground_truth=ground_truth,\n            neurons=neurons,\n            mode=\"suppress\",\n            undo_modification=undo_modification,\n            quiet=quiet,\n        )\n\n    def enhance_knowledge(\n        self,\n        prompt: str,\n        ground_truth: str,\n        neurons: List[List[int]],\n        undo_modification: bool = True,\n        quiet: bool = False,\n    ) -> Tuple[dict, Callable]:\n        \"\"\"\n        prompt the model with `prompt`, multiplying the activations at the positions\n        specified by `neurons` by 2, and measure the resulting affect on the ground truth probability.\n        \"\"\"\n        return self.modify_activations(\n            prompt=prompt,\n            ground_truth=ground_truth,\n            neurons=neurons,\n            mode=\"enhance\",\n            undo_modification=undo_modification,\n            quiet=quiet,\n        )\n\n    @torch.no_grad()\n    def modify_weights(\n        self,\n        prompt: str,\n        neurons: List[List[int]],\n        target: str,\n        mode: str = \"edit\",\n        erase_value: str = \"zero\",\n        undo_modification: bool = True,\n        quiet: bool = False,\n    ) -> Tuple[dict, Callable]:\n        \"\"\"\n        Update the *weights* of the neural net in the positions specified by `neurons`.\n        Specifically, the weights of the second Linear layer in the ff are updated by adding or subtracting the value\n        of the word embeddings for `target`.\n        \"\"\"\n        assert mode in [\"edit\", \"erase\"]\n        assert erase_value in [\"zero\", \"unk\"]\n        results_dict = {}\n\n        _, _, target_label = self._prepare_inputs(prompt, target)\n        # get the baseline probabilities of the target being generated + the argmax / greedy completion before modifying the weights\n        (\n            gt_baseline_prob,\n            argmax_baseline_prob,\n            argmax_completion_str,\n            argmax_tokens,\n        ) = self._generate(prompt, target)\n        if not quiet:\n            print(\n                f\"\\nBefore modification - groundtruth probability: {gt_baseline_prob}\\nArgmax completion: `{argmax_completion_str}`\\nArgmax prob: {argmax_baseline_prob}\"\n            )\n        results_dict[\"before\"] = {\n            \"gt_prob\": gt_baseline_prob,\n            \"argmax_completion\": argmax_completion_str,\n            \"argmax_prob\": argmax_baseline_prob,\n        }\n\n        # get the word embedding values of the baseline + target predictions\n        word_embeddings_weights = self._get_word_embeddings()\n        if mode == \"edit\":\n            assert (\n                self.model_type == \"bert\"\n            ), \"edit mode currently only working for bert models - TODO\"\n            original_prediction_id = argmax_tokens[0]\n            original_prediction_embedding = word_embeddings_weights[\n                original_prediction_id\n            ]\n            target_embedding = word_embeddings_weights[target_label]\n\n        if erase_value == \"zero\":\n            erase_value = 0\n        else:\n            assert self.model_type == \"bert\", \"GPT models don't have an unk token\"\n            erase_value = word_embeddings_weights[self.unk_token]\n\n        # modify the weights by subtracting the original prediction's word embedding\n        # and adding the target embedding\n        original_weight_values = []  # to reverse the action later\n        for layer_idx, position in neurons:\n            output_ff_weights = self._get_output_ff_layer(layer_idx)\n            if self.model_type == \"gpt2\":\n                # since gpt2 uses a conv1d layer instead of a linear layer in the ff block, the weights are in a different format\n                original_weight_values.append(\n                    output_ff_weights[position, :].detach().clone()\n                )\n            else:\n                original_weight_values.append(\n                    output_ff_weights[:, position].detach().clone()\n                )\n            if mode == \"edit\":\n                if self.model_type == \"gpt2\":\n                    output_ff_weights[position, :] -= original_prediction_embedding * 2\n                    output_ff_weights[position, :] += target_embedding * 2\n                else:\n                    output_ff_weights[:, position] -= original_prediction_embedding * 2\n                    output_ff_weights[:, position] += target_embedding * 2\n            else:\n                if self.model_type == \"gpt2\":\n                    output_ff_weights[position, :] = erase_value\n                else:\n                    output_ff_weights[:, position] = erase_value\n\n        # get the probabilities of the target being generated + the argmax / greedy completion after modifying the weights\n        (\n            new_gt_prob,\n            new_argmax_prob,\n            new_argmax_completion_str,\n            new_argmax_tokens,\n        ) = self._generate(prompt, target)\n        if not quiet:\n            print(\n                f\"\\nAfter modification - groundtruth probability: {new_gt_prob}\\nArgmax completion: `{new_argmax_completion_str}`\\nArgmax prob: {new_argmax_prob}\"\n            )\n        results_dict[\"after\"] = {\n            \"gt_prob\": new_gt_prob,\n            \"argmax_completion\": new_argmax_completion_str,\n            \"argmax_prob\": new_argmax_prob,\n        }\n\n        def unpatch_fn():\n            # reverse modified weights\n            for idx, (layer_idx, position) in enumerate(neurons):\n                output_ff_weights = self._get_output_ff_layer(layer_idx)\n                if self.model_type == \"gpt2\":\n                    output_ff_weights[position, :] = original_weight_values[idx]\n                else:\n                    output_ff_weights[:, position] = original_weight_values[idx]\n\n        if undo_modification:\n            unpatch_fn()\n            unpatch_fn = lambda *args: args\n\n        return results_dict, unpatch_fn\n\n    def edit_knowledge(\n        self,\n        prompt: str,\n        target: str,\n        neurons: List[List[int]],\n        undo_modification: bool = True,\n        quiet: bool = False,\n    ) -> Tuple[dict, Callable]:\n        return self.modify_weights(\n            prompt=prompt,\n            neurons=neurons,\n            target=target,\n            mode=\"edit\",\n            undo_modification=undo_modification,\n            quiet=quiet,\n        )\n\n    def erase_knowledge(\n        self,\n        prompt: str,\n        neurons: List[List[int]],\n        erase_value: str = \"zero\",\n        target: Optional[str] = None,\n        undo_modification: bool = True,\n        quiet: bool = False,\n    ) -> Tuple[dict, Callable]:\n        return self.modify_weights(\n            prompt=prompt,\n            neurons=neurons,\n            target=target,\n            mode=\"erase\",\n            erase_value=erase_value,\n            undo_modification=undo_modification,\n            quiet=quiet,\n        )","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:30:54.657621Z","iopub.execute_input":"2022-11-28T16:30:54.658036Z","iopub.status.idle":"2022-11-28T16:30:54.754147Z","shell.execute_reply.started":"2022-11-28T16:30:54.658003Z","shell.execute_reply":"2022-11-28T16:30:54.753127Z"},"trusted":true},"execution_count":15,"outputs":[]},{"cell_type":"code","source":"\n\ndef get_attributes(x: nn.Module, attributes: str):\n    \"\"\"\n    gets a list of period-separated attributes\n    i.e get_attributes(model, 'transformer.encoder.layer')\n        should return the same as model.transformer.encoder.layer\n    \"\"\"\n    for attr in attributes.split(\".\"):\n        x = getattr(x, attr)\n    return x\n\n\ndef set_attribute_recursive(x: nn.Module, attributes: \"str\", new_attribute: nn.Module):\n    \"\"\"\n    Given a list of period-separated attributes - set the final attribute in that list to the new value\n    i.e set_attribute_recursive(model, 'transformer.encoder.layer', NewLayer)\n        should set the final attribute of model.transformer.encoder.layer to NewLayer\n    \"\"\"\n    for attr in attributes.split(\".\")[:-1]:\n        x = getattr(x, attr)\n    setattr(x, attributes.split(\".\")[-1], new_attribute)\n\n\ndef get_ff_layer(\n    model: nn.Module,\n    layer_idx: int,\n    transformer_layers_attr: str = \"bert.encoder.layer\",\n    ff_attrs: str = \"intermediate\",\n):\n    \"\"\"\n    Gets the feedforward layer of a model within the transformer block\n    `model`: torch.nn.Module\n      a torch.nn.Module\n    `layer_idx`: int\n      which transformer layer to access\n    `transformer_layers_attr`: str\n      chain of attributes (separated by periods) that access the transformer layers within `model`.\n      The transformer layers are expected to be indexable - i.e a Modulelist\n    `ff_attrs`: str\n      chain of attributes (separated by periods) that access the ff block within a transformer layer\n    \"\"\"\n    transformer_layers = get_attributes(model, transformer_layers_attr)\n    assert layer_idx < len(\n        transformer_layers\n    ), f\"cannot get layer {layer_idx + 1} of a {len(transformer_layers)} layer model\"\n    ff_layer = get_attributes(transformer_layers[layer_idx], ff_attrs)\n    return ff_layer\n\n\ndef register_hook(\n    model: nn.Module,\n    layer_idx: int,\n    f: Callable,\n    transformer_layers_attr: str = \"bert.encoder.layer\",\n    ff_attrs: str = \"intermediate\",\n):\n    \"\"\"\n    Registers a forward hook in a pytorch transformer model that applies some function, f, to the intermediate\n    activations of the transformer model.\n\n    specify how to access the transformer layers (which are expected to be indexable - i.e a ModuleList) with transformer_layers_attr\n    and how to access the ff layer with ff_attrs\n\n    `model`: torch.nn.Module\n      a torch.nn.Module\n    `layer_idx`: int\n      which transformer layer to access\n    `f`: Callable\n      a callable function that takes in the intermediate activations\n    `transformer_layers_attr`: str\n      chain of attributes (separated by periods) that access the transformer layers within `model`.\n      The transformer layers are expected to be indexable - i.e a Modulelist\n    `ff_attrs`: str\n      chain of attributes (separated by periods) that access the ff block within a transformer layer\n    \"\"\"\n    ff_layer = get_ff_layer(\n        model,\n        layer_idx,\n        transformer_layers_attr=transformer_layers_attr,\n        ff_attrs=ff_attrs,\n    )\n\n    def hook_fn(m, i, o):\n        f(o)\n\n    return ff_layer.register_forward_hook(hook_fn)\n\n\nclass Patch(torch.nn.Module):\n    \"\"\"\n    Patches a torch module to replace/suppress/enhance the intermediate activations\n    \"\"\"\n\n    def __init__(\n        self,\n        ff_layer: nn.Module,\n        mask_idx: int,\n        replacement_activations: torch.Tensor = None,\n        target_positions: List[List[int]] = None,\n        mode: str = \"replace\",\n        enhance_value: float = 2.0,\n    ):\n        super().__init__()\n        self.ff = ff_layer\n        self.acts = replacement_activations\n        self.mask_idx = mask_idx\n        self.target_positions = target_positions\n        self.enhance_value = enhance_value\n        assert mode in [\"replace\", \"suppress\", \"enhance\"]\n        self.mode = mode\n        if self.mode == \"replace\":\n            assert self.acts is not None\n        elif self.mode in [\"enhance\", \"suppress\"]:\n            assert self.target_positions is not None\n\n    def forward(self, x: torch.Tensor):\n        x = self.ff(x)\n        if self.mode == \"replace\":\n            x[:, self.mask_idx, :] = self.acts\n        elif self.mode == \"suppress\":\n            for pos in self.target_positions:\n                x[:, self.mask_idx, pos] = 0.0\n        elif self.mode == \"enhance\":\n            for pos in self.target_positions:\n                x[:, self.mask_idx, pos] *= self.enhance_value\n        else:\n            raise NotImplementedError\n        return x\n\n\ndef patch_ff_layer(\n    model: nn.Module,\n    mask_idx: int,\n    layer_idx: int = None,\n    replacement_activations: torch.Tensor = None,\n    mode: str = \"replace\",\n    transformer_layers_attr: str = \"bert.encoder.layer\",\n    ff_attrs: str = \"intermediate\",\n    neurons: List[List[int]] = None,\n):\n    \"\"\"\n    replaces the ff layer at `layer_idx` with a `Patch` class - that will replace the intermediate activations at sequence position\n    `mask_index` with `replacement_activations`\n\n    `model`: nn.Module\n      a torch.nn.Module [currently only works with HF Bert models]\n    `layer_idx`: int\n      which transformer layer to access\n    `mask_idx`: int\n      the index (along the sequence length) of the activation to replace.\n      TODO: multiple indices\n    `replacement_activations`: torch.Tensor\n      activations [taken from the mask_idx position of the unmodified activations] of shape [b, d]\n    `transformer_layers_attr`: str\n      chain of attributes (separated by periods) that access the transformer layers within `model`.\n      The transformer layers are expected to be indexable - i.e a Modulelist\n    `ff_attrs`: str\n      chain of attributes (separated by periods) that access the ff block within a transformer layer\n    \"\"\"\n    transformer_layers = get_attributes(model, transformer_layers_attr)\n\n    if mode == \"replace\":\n        ff_layer = get_attributes(transformer_layers[layer_idx], ff_attrs)\n        assert layer_idx < len(\n            transformer_layers\n        ), f\"cannot get layer {layer_idx + 1} of a {len(transformer_layers)} layer model\"\n\n        set_attribute_recursive(\n            transformer_layers[layer_idx],\n            ff_attrs,\n            Patch(\n                ff_layer,\n                mask_idx,\n                replacement_activations=replacement_activations,\n                mode=mode,\n            ),\n        )\n\n    elif mode in [\"suppress\", \"enhance\"]:\n        neurons_dict = collections.defaultdict(list)\n        for neuron in neurons:\n            layer_idx, pos = neuron\n            neurons_dict[layer_idx].append(pos)\n        for layer_idx, positions in neurons_dict.items():\n            assert layer_idx < len(transformer_layers)\n            ff_layer = get_attributes(transformer_layers[layer_idx], ff_attrs)\n            set_attribute_recursive(\n                transformer_layers[layer_idx],\n                ff_attrs,\n                Patch(\n                    ff_layer,\n                    mask_idx,\n                    replacement_activations=None,\n                    mode=mode,\n                    target_positions=positions,\n                ),\n            )\n    else:\n        raise NotImplementedError\n\n\ndef unpatch_ff_layer(\n    model: nn.Module,\n    layer_idx: int,\n    transformer_layers_attr: str = \"bert.encoder.layer\",\n    ff_attrs: str = \"intermediate\",\n):\n    \"\"\"\n    Removes the `Patch` applied by `patch_ff_layer`, replacing it with its original value.\n\n    `model`: torch.nn.Module\n      a torch.nn.Module [currently only works with HF Bert models]\n    `layer_idx`: int\n      which transformer layer to access\n    `transformer_layers_attr`: str\n      chain of attributes (separated by periods) that access the transformer layers within `model`.\n      The transformer layers are expected to be indexable - i.e a Modulelist\n    `ff_attrs`: str\n      chain of attributes (separated by periods) that access the ff block within a transformer layer\n    \"\"\"\n    transformer_layers = get_attributes(model, transformer_layers_attr)\n    assert layer_idx < len(\n        transformer_layers\n    ), f\"cannot get layer {layer_idx + 1} of a {len(transformer_layers)} layer model\"\n    ff_layer = get_attributes(transformer_layers[layer_idx], ff_attrs)\n    assert isinstance(ff_layer, Patch), \"Can't unpatch a layer that hasn't been patched\"\n    set_attribute_recursive(\n        transformer_layers[layer_idx],\n        ff_attrs,\n        ff_layer.ff,\n    )\n\n\ndef unpatch_ff_layers(\n    model: nn.Module,\n    layer_indices: int,\n    transformer_layers_attr: str = \"bert.encoder.layer\",\n    ff_attrs: str = \"intermediate\",\n):\n    \"\"\"\n    Calls unpatch_ff_layer for all layers in layer_indices\n    \"\"\"\n    for layer_idx in layer_indices:\n        unpatch_ff_layer(model, layer_idx, transformer_layers_attr, ff_attrs)","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:03:08.541509Z","iopub.execute_input":"2022-11-28T16:03:08.542152Z","iopub.status.idle":"2022-11-28T16:03:08.567054Z","shell.execute_reply.started":"2022-11-28T16:03:08.542115Z","shell.execute_reply":"2022-11-28T16:03:08.565926Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"code","source":"RESULTS_DIR = Path('results')\nos.makedirs(RESULTS_DIR, exist_ok=True)\nrandom.seed(2022)\n\nMODEL_NAME= \"bert-base-uncased\"\n# these are some hyperparameters for the integrated gradients step\nBATCH_SIZE = 20\nSTEPS = 20 # number of steps in the integrated grad calculation\nADAPTIVE_THRESHOLD = 0.3 # in the paper, they find the threshold value `t` by multiplying the max attribution score by some float - this is that float.\nP = 0.5 # the threshold for the sharing percentage\n\nmodel, tokenizer = initialize_model_and_tokenizer(MODEL_NAME)","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:03:08.568619Z","iopub.execute_input":"2022-11-28T16:03:08.569376Z","iopub.status.idle":"2022-11-28T16:04:03.226654Z","shell.execute_reply.started":"2022-11-28T16:03:08.569337Z","shell.execute_reply":"2022-11-28T16:04:03.225642Z"},"trusted":true},"execution_count":6,"outputs":[{"output_type":"display_data","data":{"text/plain":"Downloading:   0%|          | 0.00/226k [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"f56c506df35e4d628e83bc18db6f88d9"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"Downloading:   0%|          | 0.00/28.0 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"d1a771df3818472995d6b36562402db9"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"Downloading:   0%|          | 0.00/570 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"ea8831a4238840aca7d7d90a81c1e495"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"Downloading:   0%|          | 0.00/420M [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"7cc07ef94d3b4a5bae3a8a98ed727b39"}},"metadata":{}},{"name":"stderr","text":"If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`\nSome weights of the model checkpoint at bert-base-uncased were not used when initializing BertLMHeadModel: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n- This IS expected if you are initializing BertLMHeadModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n- This IS NOT expected if you are initializing BertLMHeadModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n","output_type":"stream"}]},{"cell_type":"code","source":"sentences= {'.': [], ',': [], '!': [], '?': [], ';': []}\nwith open('/kaggle/input/english-dataset-for-punctuation-prediction/English-OPUS.txt', 'r') as f:\n    for line in f.readlines():\n        line= line.strip('\\n').strip(' ')\n        if len(line)!= 0 and line[-1] in sentences:\n            if len(sentences[line[-1]])== 200:\n                continue\n            sentences[line[-1]].append(line[:-1]+ ' [MASK]')","metadata":{"execution":{"iopub.status.busy":"2022-12-13T01:27:24.061131Z","iopub.execute_input":"2022-12-13T01:27:24.061629Z","iopub.status.idle":"2022-12-13T01:27:25.977284Z","shell.execute_reply.started":"2022-12-13T01:27:24.061583Z","shell.execute_reply":"2022-12-13T01:27:25.975968Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"code","source":"sentences","metadata":{"execution":{"iopub.status.busy":"2022-12-13T01:31:34.530598Z","iopub.execute_input":"2022-12-13T01:31:34.531063Z","iopub.status.idle":"2022-12-13T01:31:34.541587Z","shell.execute_reply.started":"2022-12-13T01:31:34.531029Z","shell.execute_reply":"2022-12-13T01:31:34.540315Z"},"trusted":true},"execution_count":6,"outputs":[{"execution_count":6,"output_type":"execute_result","data":{"text/plain":"{'.': 685608, ',': 6770, '!': 56821, '?': 119442, ';': 1166}"},"metadata":{}}]},{"cell_type":"code","source":"sentences= {'.': 0, ',': 0, '!': 0, '?': 0, ';': 0}\nwith open('/kaggle/input/english-dataset-for-punctuation-prediction/English-OPUS.txt', 'r') as f:\n     for line in f.readlines():\n        line= line.strip('\\n').strip(' ')\n        if len(line)!= 0 and line[-1] in sentences:\n            sentences[line[-1]]+= 1","metadata":{"execution":{"iopub.status.busy":"2022-12-13T01:30:28.310694Z","iopub.execute_input":"2022-12-13T01:30:28.311996Z","iopub.status.idle":"2022-12-13T01:30:29.439908Z","shell.execute_reply.started":"2022-12-13T01:30:28.311931Z","shell.execute_reply":"2022-12-13T01:30:29.437962Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"code","source":"sentences['.'][0]","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:53:56.266702Z","iopub.execute_input":"2022-11-28T16:53:56.267222Z","iopub.status.idle":"2022-11-28T16:53:56.279990Z","shell.execute_reply.started":"2022-11-28T16:53:56.267178Z","shell.execute_reply":"2022-11-28T16:53:56.279119Z"},"trusted":true},"execution_count":23,"outputs":[{"execution_count":23,"output_type":"execute_result","data":{"text/plain":"'Today is the International Mother Language Day, an annual event in UNESCO member states to promote linguistic and cultural diversity and multilingualism [MASK]'"},"metadata":{}}]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"kn = KnowledgeNeurons(model, tokenizer, model_type=model_type(MODEL_NAME))\n# use the integrated gradients technique to find some refined neurons for your set of prompts\nrefined_neurons = kn.get_refined_neurons(\n    sentences['.'],\n    '.',\n    p= 1.0,\n    batch_size=BATCH_SIZE,\n    steps=STEPS,\n    coarse_adaptive_threshold=ADAPTIVE_THRESHOLD,\n)","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:31:01.557919Z","iopub.execute_input":"2022-11-28T16:31:01.558841Z","iopub.status.idle":"2022-11-28T16:33:31.707461Z","shell.execute_reply.started":"2022-11-28T16:31:01.558804Z","shell.execute_reply":"2022-11-28T16:33:31.706433Z"},"trusted":true},"execution_count":17,"outputs":[{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [02:30<00:00,  1.33it/s]","output_type":"stream"},{"name":"stdout","text":"\n953 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"\n","output_type":"stream"}]},{"cell_type":"code","source":"refined_neurons","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:41:18.958550Z","iopub.execute_input":"2022-11-28T16:41:18.959135Z","iopub.status.idle":"2022-11-28T16:41:18.967916Z","shell.execute_reply.started":"2022-11-28T16:41:18.959062Z","shell.execute_reply":"2022-11-28T16:41:18.966841Z"},"trusted":true},"execution_count":20,"outputs":[{"execution_count":20,"output_type":"execute_result","data":{"text/plain":"{0.6: [[8, 839], [8, 1932], [9, 927], [10, 1458]],\n 0.7: [[8, 839], [8, 1932], [9, 927], [10, 1458]],\n 0.8: [[8, 839], [9, 927], [10, 1458]],\n 0.9: [[9, 927]],\n 1.0: []}"},"metadata":{}}]},{"cell_type":"code","source":"neurons= {'.': [], ',': [], '!': [], '?': [], ';': []}\nfor key in sentences:\n    texts= sentences[key]\n    refined_neurons = kn.get_refined_neurons(\n        texts,\n        key,\n        p= 1.0,\n        batch_size=BATCH_SIZE,\n        steps=STEPS,\n        coarse_adaptive_threshold=ADAPTIVE_THRESHOLD,\n    )\n    neurons[key].append(refined_neurons)","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:41:23.698635Z","iopub.execute_input":"2022-11-28T16:41:23.699005Z","iopub.status.idle":"2022-11-28T16:51:18.982742Z","shell.execute_reply.started":"2022-11-28T16:41:23.698975Z","shell.execute_reply":"2022-11-28T16:51:18.981754Z"},"trusted":true},"execution_count":21,"outputs":[{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [02:29<00:00,  1.34it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n953 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [01:34<00:00,  2.12it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n903 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [02:08<00:00,  1.55it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n1809 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [01:49<00:00,  1.82it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n953 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [01:52<00:00,  1.77it/s]","output_type":"stream"},{"name":"stdout","text":"\n1535 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"\n","output_type":"stream"}]},{"cell_type":"code","source":"neurons","metadata":{"execution":{"iopub.status.busy":"2022-11-28T16:52:23.219381Z","iopub.execute_input":"2022-11-28T16:52:23.219867Z","iopub.status.idle":"2022-11-28T16:52:23.235457Z","shell.execute_reply.started":"2022-11-28T16:52:23.219825Z","shell.execute_reply":"2022-11-28T16:52:23.233419Z"},"trusted":true},"execution_count":22,"outputs":[{"execution_count":22,"output_type":"execute_result","data":{"text/plain":"{'.': [{0.6: [[8, 839], [8, 1932], [9, 927], [10, 1458]],\n   0.7: [[8, 839], [8, 1932], [9, 927], [10, 1458]],\n   0.8: [[8, 839], [9, 927], [10, 1458]],\n   0.9: [[9, 927]],\n   1.0: []}],\n ',': [{0.6: [[11, 1205], [11, 2366], [11, 240]],\n   0.7: [[11, 1205], [11, 2366]],\n   0.8: [[11, 1205]],\n   0.9: [[11, 1205]],\n   1.0: []}],\n '!': [{0.6: [[10, 431], [11, 240], [11, 2399]],\n   0.7: [],\n   0.8: [],\n   0.9: [],\n   1.0: []}],\n '?': [{0.6: [[9, 1673], [10, 861]],\n   0.7: [[9, 1673]],\n   0.8: [[9, 1673]],\n   0.9: [],\n   1.0: []}],\n ';': [{0.6: [[11, 240]],\n   0.7: [[11, 240]],\n   0.8: [[11, 240]],\n   0.9: [],\n   1.0: []}]}"},"metadata":{}}]},{"cell_type":"code","source":"gpt, gpt_tokenizer = initialize_model_and_tokenizer('gpt2')","metadata":{"execution":{"iopub.status.busy":"2022-11-28T17:36:21.114207Z","iopub.execute_input":"2022-11-28T17:36:21.114584Z","iopub.status.idle":"2022-11-28T17:36:31.984020Z","shell.execute_reply.started":"2022-11-28T17:36:21.114553Z","shell.execute_reply":"2022-11-28T17:36:31.983003Z"},"trusted":true},"execution_count":27,"outputs":[]},{"cell_type":"code","source":"gpt_kn = KnowledgeNeurons(gpt, gpt_tokenizer, model_type=model_type('gpt2'))","metadata":{"execution":{"iopub.status.busy":"2022-11-28T17:36:31.989082Z","iopub.execute_input":"2022-11-28T17:36:31.991280Z","iopub.status.idle":"2022-11-28T17:36:32.145616Z","shell.execute_reply.started":"2022-11-28T17:36:31.991241Z","shell.execute_reply":"2022-11-28T17:36:32.144554Z"},"trusted":true},"execution_count":28,"outputs":[]},{"cell_type":"code","source":"neurons= {'.': [], ',': [], '!': [], '?': [], ';': []}\nfor key in sentences:\n    texts= sentences[key]\n    refined_neurons = gpt_kn.get_refined_neurons(\n        texts,\n        key,\n        p= 1.0,\n        batch_size=BATCH_SIZE,\n        steps=STEPS,\n        coarse_adaptive_threshold=ADAPTIVE_THRESHOLD,\n    )\n    neurons[key].append(refined_neurons)","metadata":{"execution":{"iopub.status.busy":"2022-11-28T17:36:32.150343Z","iopub.execute_input":"2022-11-28T17:36:32.152605Z","iopub.status.idle":"2022-11-28T17:47:28.770353Z","shell.execute_reply.started":"2022-11-28T17:36:32.152564Z","shell.execute_reply":"2022-11-28T17:47:28.769362Z"},"trusted":true},"execution_count":29,"outputs":[{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [02:49<00:00,  1.18it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n429 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [01:41<00:00,  1.96it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n854 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [02:22<00:00,  1.40it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n811 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [01:59<00:00,  1.68it/s]\n","output_type":"stream"},{"name":"stdout","text":"\n863 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"Getting coarse neurons for each prompt...: 100%|██████████| 200/200 [02:03<00:00,  1.62it/s]","output_type":"stream"},{"name":"stdout","text":"\n956 coarse neurons found - refining\n5 neurons remaining after refining\n","output_type":"stream"},{"name":"stderr","text":"\n","output_type":"stream"}]},{"cell_type":"code","source":"neurons","metadata":{"execution":{"iopub.status.busy":"2022-11-28T17:47:46.368158Z","iopub.execute_input":"2022-11-28T17:47:46.368527Z","iopub.status.idle":"2022-11-28T17:47:46.376622Z","shell.execute_reply.started":"2022-11-28T17:47:46.368497Z","shell.execute_reply":"2022-11-28T17:47:46.375640Z"},"trusted":true},"execution_count":30,"outputs":[{"execution_count":30,"output_type":"execute_result","data":{"text/plain":"{'.': [{0.6: [[0, 990], [9, 831]],\n   0.7: [[9, 831]],\n   0.8: [[9, 831]],\n   0.9: [[9, 831]],\n   1.0: []}],\n ',': [{0.6: [[9, 831]],\n   0.7: [[9, 831]],\n   0.8: [[9, 831]],\n   0.9: [[9, 831]],\n   1.0: []}],\n '!': [{0.6: [[9, 831], [0, 990]],\n   0.7: [[9, 831]],\n   0.8: [[9, 831]],\n   0.9: [[9, 831]],\n   1.0: []}],\n '?': [{0.6: [[0, 990], [10, 897], [9, 831]],\n   0.7: [[0, 990], [10, 897], [9, 831]],\n   0.8: [[9, 831]],\n   0.9: [],\n   1.0: []}],\n ';': [{0.6: [[9, 831]],\n   0.7: [[9, 831]],\n   0.8: [[9, 831]],\n   0.9: [[9, 831]],\n   1.0: []}]}"},"metadata":{}}]}]}