TCC(def).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "7KjzUw0K-49l",
        "EbXloDtnTyLp",
        "QTiTogfFLEIc",
        "modtyyuJnuAj",
        "KiDJkduU-Tr6",
        "bPHEzd3qJba0",
        "W1W0sGoQCL-j",
        "06akBZ3zC1S8",
        "MWWUVgJFApw9",
        "JGUDy8-MHPSb",
        "Prgx47fxTC8L",
        "Z0MMG3AJCJIM",
        "-_QbTgYyGpZt",
        "Kg877UW39n9p",
        "t74soB3f3QtD",
        "5RG3I0WD3nKS",
        "G-nYzzaINwbq",
        "6WIel2kEUaeM",
        "JGLOTDg2FxlB",
        "5n03ggfcKxhN",
        "k_wr704MxE1R",
        "ElMEXgzTB1nJ",
        "0-Io8O8IzzxH",
        "gljQcB8F2ksn",
        "hDqjb6E0EyX2",
        "iBdeOEdspqbO",
        "wzKZLaERrJ5g",
        "afnh5NRykrJL",
        "isCeG4KH6hqn",
        "dPUzUls7TWcu"
      ],
      "gpuType": "T4",
      "private_outputs": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# TODO\n",
        "\n",
        "\n",
        "constractive destilattion based on softmax predicts of cosine similarity and a distanece\n",
        "$=\\exp{[ s(p_T,p_S])}/\\sum_i\\exp{[s(p^{(i)}_T,p^{(i)}_S)]}$\n",
        "\n",
        "\n",
        "pipelines\n",
        "- first investigation\n",
        "  - diferent modelsizes\n",
        "- optimization\n",
        "  - sparse, quantize, distillation, (sparse + quantize), (sparse + quantize+ lora), (distillation + sparse + quantize+ lora)\n",
        "\n",
        "- data aumentation based in the ranked\n",
        "- apply finds to all datasets\n",
        "  - select datasets\n",
        "- write saves in images to drive\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "7KjzUw0K-49l"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install megatron-core\n",
        "\n",
        "\n",
        "# https://github.com/NVIDIA/Megatron-LM/commit/555036e88fa7129e7caaf67b0c1ec838703c35dd#diff-3a709bfafbc79eb07662cc059fadc6b8cca7aa6083f16831fd029865636bf298\n",
        "# https://github.com/huggingface/peft/blob/bf54136a79cc85b0e4c3915b4e1eb158f43c4b73/src/peft/tuners/lora/tp_layer.py#L26\n",
        "\n",
        "# from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear\n",
        "\n",
        "# core.transformer.transformer_config."
      ],
      "metadata": {
        "id": "j4xsjye5HG2w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# cuda torch.copile() aceleration\n",
        "\n",
        "\n",
        "# !export LC_ALL=\"en_US.UTF-8\"\n",
        "# !export LD_LIBRARY_PATH=\"/usr/lib64-nvidia\"\n",
        "# !export LIBRARY_PATH=\"/usr/local/cuda/lib64/stubs\"\n",
        "# !ldconfig /usr/lib64-nvidia"
      ],
      "metadata": {
        "id": "YX25sAW7w83n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Init"
      ],
      "metadata": {
        "id": "EbXloDtnTyLp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "dKo3j5CDJq4d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "path = \"/content/gdrive/MyDrive/Projetos/TCC/slr_data_tcc.zip\""
      ],
      "metadata": {
        "id": "WwPkOYTfKBa2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -O data.zip https://www.dropbox.com/sh/or0eyfo8znyu2kp/AABxXJVII48U0vY8TT3Bbp6Ea?dl=0\n",
        "!unzip \"./gdrive/MyDrive/Projetos/TCC/slr_data_tcc.zip\"\n",
        "\n",
        "# !unzip OriData/DATA_SLR.zip"
      ],
      "metadata": {
        "id": "NLVsqxRTBWC6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install -U adapter-transformers peft\n",
        "# !pip install -q peft\n",
        "!pip install -q einops einop\n",
        "!pip install -q bitsandbytes accelerate datasets\n",
        "!pip install -q git+https://github.com/huggingface/transformers.git@main git+https://github.com/huggingface/peft.git\n",
        "!pip install -q git+https://github.com/huggingface/accelerate.git@main\n",
        "\n",
        "# !pip install -U adapter-transformers"
      ],
      "metadata": {
        "id": "_ILbl7NeckEa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Imports\n"
      ],
      "metadata": {
        "id": "QTiTogfFLEIc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import argparse\n",
        "import json\n",
        "import logging\n",
        "import math\n",
        "import os\n",
        "import random\n",
        "from pathlib import Path\n",
        "from tqdm import tqdm\n",
        "\n",
        "# import datasets\n",
        "# from datasets import load_dataset, DatasetDict\n",
        "\n",
        "# import evaluate\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import transformers\n",
        "\n",
        "# from google.colab import userdata\n",
        "# userdata.get('HF_TOKEN')"
      ],
      "metadata": {
        "id": "39Odq6GoLF9_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##  Utils functions"
      ],
      "metadata": {
        "id": "modtyyuJnuAj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "def read_json(file_path):\n",
        "    \"\"\"\n",
        "    Reads JSON data from a file.\n",
        "\n",
        "    Args:\n",
        "    - file_path (str): The path to the JSON file.\n",
        "\n",
        "    Returns:\n",
        "    - data (dict): The JSON data read from the file.\n",
        "    \"\"\"\n",
        "    try:\n",
        "        with open(file_path, 'r') as json_file:\n",
        "            data = json.load(json_file)\n",
        "        return data\n",
        "    except FileNotFoundError:\n",
        "        print(f\"Error: File '{file_path}' not found.\")\n",
        "        return None\n",
        "    except json.JSONDecodeError:\n",
        "        print(f\"Error: Invalid JSON format in file '{file_path}'.\")\n",
        "        return None\n",
        "\n",
        "\n",
        "def create_subset(original_dict, keys_to_keep):\n",
        "    \"\"\"\n",
        "    Create a subset dictionary containing only the specified keys from the original dictionary.\n",
        "\n",
        "    Args:\n",
        "    - original_dict (dict): The original dictionary.\n",
        "    - keys_to_keep (list): A list of keys to keep in the subset dictionary.\n",
        "\n",
        "    Returns:\n",
        "    - subset_dict (dict): The subset dictionary containing only the specified keys.\n",
        "    \"\"\"\n",
        "    subset_dict = {}\n",
        "    for key in keys_to_keep:\n",
        "        if key in original_dict:\n",
        "            subset_dict[key] = original_dict[key]\n",
        "    return subset_dict\n",
        "\n",
        "\n",
        "import os\n",
        "def data_import(file_path, n=8, samples_n=200, samples_p=1):\n",
        "    print(\"####\"*5)\n",
        "    print()\n",
        "    print(os.path.basename(file_path))\n",
        "    df = preprocess_csv_file(file_path)\n",
        "\n",
        "\n",
        "    tk = lambda x: len(tokenizer_t_ref.encode(x))\n",
        "    tk_len = df.text.map(tk)\n",
        "    tk_len.hist()\n",
        "\n",
        "\n",
        "    df_rest = df[tk_len > 20]\n",
        "    df_train = df_rest.groupby('label').sample(n)\n",
        "\n",
        "    print(\"data count:\")\n",
        "    print(df['label'].value_counts())\n",
        "\n",
        "    print(df_train.iloc[0])\n",
        "\n",
        "    df_test = df[~df.index.isin(df_train.index)]\n",
        "\n",
        "    if len(df_test[df_test['label'] == 1]) < samples_p:\n",
        "        df_test_one = df_test[df_test['label'] == 1].copy()\n",
        "    else:\n",
        "        df_test_one = df_test[df_test['label'] == 1].sample(n=samples_p, replace=False)\n",
        "\n",
        "\n",
        "    if len(df_test[df_test['label'] == 0]) < samples_n:\n",
        "        df_test_zero = df_test[df_test['label'] == 0].copy()\n",
        "    else:\n",
        "        df_test_zero = df_test[df_test['label'] == 0].sample(n=samples_n, replace=False)\n",
        "\n",
        "    exs_test = pd.concat([df_test_one,df_test_zero])\n",
        "\n",
        "    print(\"Eval count:\")\n",
        "    print(exs_test['label'].value_counts())\n",
        "    return df_train, exs_test\n",
        "\n",
        "\n",
        "def loaders(df_train, exs_test, config):\n",
        "    sim_data = SIMDataset(df_train, tokenize_batch_t, config[\"n_pairs\"],\n",
        "                          )\n",
        "    train_dataloader = DataLoader(sim_data,shuffle=True,\n",
        "                                  batch_size=config[\"batch_size\"])\n",
        "\n",
        "    test_dataloader_eval = DataLoader(shuffle=False, batch_size=config[\"batch_size_eval\"],\n",
        "                dataset=SIMDataset(\n",
        "                    exs_test, tokenize_batch_t, 0,\n",
        "                          )\n",
        "                )\n",
        "\n",
        "    train_dataloader_eval = DataLoader(shuffle=False, batch_size=config[\"batch_size\"],\n",
        "                dataset=SIMDataset(\n",
        "                    df_train, tokenize_batch_t, 0,\n",
        "                    # emb= emb, device=\"cuda\", pttn=\"* d a\"\n",
        "                          )\n",
        "                )\n",
        "\n",
        "    return train_dataloader, test_dataloader_eval,  train_dataloader_eval\n"
      ],
      "metadata": {
        "id": "JYAojGxGHTdf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "def create_repo_if_not_exist(repo_path):\n",
        "    directory = os.path.dirname(repo_path)\n",
        "    os.makedirs(directory, exist_ok=True)\n",
        "    print(\"created!\",repo_path)\n",
        "\n"
      ],
      "metadata": {
        "id": "_XgOcvRSgRiA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import os\n",
        "\n",
        "def count_non_zero_params(model):\n",
        "    total_params = 0\n",
        "    non_zero_params = 0\n",
        "\n",
        "    for param in model.parameters():\n",
        "        to = param.numel()\n",
        "        non = torch.count_nonzero(param).item()\n",
        "        non = 0 if non is None else non\n",
        "        # print(non)\n",
        "        # print(to)\n",
        "        # print(non/to)\n",
        "\n",
        "        total_params += to  # Count total number of parameters\n",
        "        non_zero_params += non  # Count non-zero elements\n",
        "        # print(non)\n",
        "        # print(non/to)*100\n",
        "\n",
        "\n",
        "    percentage_non_zero = (non_zero_params / total_params) * 100\n",
        "\n",
        "    return percentage_non_zero\n",
        "\n",
        "def export_model(model, file_path):\n",
        "    \"\"\"\n",
        "    Save a PyTorch model to a file.\n",
        "\n",
        "    Parameters:\n",
        "        model (state_dict): The PyTorch state_dict of the model to be saved.\n",
        "        file_path (str): Path to the file where the model will be saved.\n",
        "\n",
        "    Returns:\n",
        "        bool: True if the model is successfully saved, False otherwise.\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Create directory if it doesn't exist\n",
        "        directory = os.path.dirname(file_path)\n",
        "        os.makedirs(directory, exist_ok=True)\n",
        "\n",
        "        # Save the model\n",
        "        torch.save(model, file_path)\n",
        "        print(f\"Exported Model: {file_path}\")\n",
        "    except Exception as e:\n",
        "        print(f\"Error saving model: {e}\")\n",
        "        return False\n",
        "\n",
        "\n",
        "import json\n",
        "import os\n",
        "\n",
        "def numpy_to_list(d):\n",
        "    \"\"\"\n",
        "    Recursively convert NumPy arrays to lists in a dictionary.\n",
        "\n",
        "    Parameters:\n",
        "        d (dict): Dictionary to be processed.\n",
        "\n",
        "    Returns:\n",
        "        dict: Dictionary with NumPy arrays converted to lists.\n",
        "    \"\"\"\n",
        "    for key, value in d.items():\n",
        "        if isinstance(value, dict):\n",
        "            d[key] = numpy_to_list(value)\n",
        "        elif isinstance(value, np.ndarray):\n",
        "            d[key] = value.tolist()\n",
        "    return d\n",
        "\n",
        "def export_json(data, file_path):\n",
        "    \"\"\"\n",
        "    Export data to a JSON file.\n",
        "\n",
        "    Parameters:\n",
        "        data (dict): Data to be exported to JSON.\n",
        "        file_path (str): Path to the JSON file.\n",
        "\n",
        "    Returns:\n",
        "        bool: True if the export is successful, False otherwise.\n",
        "    \"\"\"\n",
        "    data = numpy_to_list(data)\n",
        "    try:\n",
        "        # Create directory if it doesn't exist\n",
        "        directory = os.path.dirname(file_path)\n",
        "        os.makedirs(directory, exist_ok=True)\n",
        "\n",
        "        # Export data to JSON file\n",
        "        with open(file_path, 'w') as json_file:\n",
        "            json.dump(data, json_file, indent=2)\n",
        "        print(f\"Exported to JSON: {file_path}\")\n",
        "    except Exception as e:\n",
        "        print(f\"Error exporting to JSON: {e}\")\n",
        "        # return False\n",
        "\n"
      ],
      "metadata": {
        "id": "a9PHuln2YoIp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "\n",
        "def preprocess_csv_file(file_path):\n",
        "    \"\"\"Preprocess CSV file.\"\"\"\n",
        "    # Carregar o arquivo CSV em um DataFrame do pandas\n",
        "    df = pd.read_csv(file_path)\n",
        "\n",
        "    # # Exibir as primeiras linhas do DataFrame\n",
        "    # df.head(2)\n",
        "\n",
        "    df = df.rename(columns={\n",
        "        \"label_included\": \"label\",\n",
        "        \"labels\": \"label\"\n",
        "        })\n",
        "\n",
        "    col = \"title\"\n",
        "    df.loc[df[col].isna(), col] = \"\"\n",
        "    col = \"abstract\"\n",
        "    df.loc[df[col].isna(), col] = \"\"\n",
        "\n",
        "    df['text'] = df['title'] + \"; \" + df['abstract']\n",
        "\n",
        "\n",
        "    return df\n"
      ],
      "metadata": {
        "id": "ZQfPjvjYTGVW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import DataLoader\n",
        "from einops import pack, unpack\n",
        "\n",
        "n = 50\n",
        "ex = 2\n",
        "input_data = torch.randn(n,ex, 5)\n",
        "output_data = torch.randn(n, 1)\n",
        "batch = 20\n",
        "\n",
        "\n",
        "def batch_load(input, output = None, batch=32):\n",
        "  input_data = input.get('data')\n",
        "  output_data = output.get('data')\n",
        "\n",
        "  input_cfg = input.get('config')\n",
        "  output_cfg = output.get('config')\n",
        "\n",
        "\n",
        "  full = input_data.shape[0]//batch\n",
        "  batch_as = [[batch]]*full\n",
        "  rest =  input_data.shape[0]%batch\n",
        "\n",
        "  if rest > 0:\n",
        "    batch_as += [[rest]]\n",
        "\n",
        "  # input_data.shape\n",
        "  in_a = unpack(input_data, batch_as, input_cfg)\n",
        "  out_a = unpack(output_data, batch_as, output_cfg)\n",
        "\n",
        "  return zip(in_a, out_a)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "a = batch_load(\n",
        "    input = {\n",
        "        \"data\": input_data,\n",
        "        \"config\": \"* ex x\"\n",
        "    },\n",
        "    output = {\n",
        "        \"data\": output_data,\n",
        "        \"config\": \"* y\"\n",
        "    },\n",
        "    batch = batch\n",
        ")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "b = iter(a)\n",
        "x = next(b)"
      ],
      "metadata": {
        "id": "z40LoWd0-sz3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Mods"
      ],
      "metadata": {
        "id": "KiDJkduU-Tr6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "todo:\n",
        "fazer a interação por loop para compartilhar os pesos junto ao copiar a arquitetra, ao inves de copiar o modelo todo de uma vez"
      ],
      "metadata": {
        "id": "bg42br3qrOII"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 📣 ModUtils"
      ],
      "metadata": {
        "id": "bPHEzd3qJba0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Main class"
      ],
      "metadata": {
        "id": "sKzzPtskCl1u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "import torch\n",
        "import copy\n",
        "from collections import OrderedDict\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "import re\n",
        "\n",
        "\n",
        "####### ajustar call para conseguir utilizar do modular depois do map\n",
        "\n",
        "\n",
        "class ModularUtils():\n",
        "    # Modular = Lora\n",
        "    # type_Mod = nn.Linear\n",
        "    # config=None\n",
        "    # class_mod = ModularUtils_modules\n",
        "\n",
        "    def __call__(self, module):\n",
        "\n",
        "      ModularUtils(\n",
        "                 module = module,\n",
        "                 Modular=self.Modular,\n",
        "                 type_Mod=self.type_Mod,\n",
        "                 config=self.config,\n",
        "                 tag=self.tag)\n",
        "\n",
        "\n",
        "      print(f'[{self.tag}.util] Done!')\n",
        "\n",
        "\n",
        "    def __init__(self,\n",
        "                 module = None,\n",
        "                 Modular=None,\n",
        "                 type_Mod=nn.Linear,\n",
        "                 config=None,\n",
        "                 tag=''\n",
        "                 ):\n",
        "      \"\"\"\n",
        "      Initialize ModularUtils.\n",
        "\n",
        "      Args:\n",
        "          Modular (callable): The custom module to replace layers.\n",
        "          type_Mod (type): The type of module to replace.\n",
        "          config (dict): Configuration for the custom module.\n",
        "      \"\"\"\n",
        "\n",
        "      self.Modular = Modular\n",
        "      self.type_Mod = type_Mod\n",
        "      self.config=config\n",
        "      self.tag=tag\n",
        "      self.module=module\n",
        "\n",
        "      if module is not None:\n",
        "        module.util = self\n",
        "\n",
        "\n",
        "    # def sub_all_linear(self, module,sub_linear=Linear4bit, **kargs):\n",
        "    #     module = self.module if (module is None) else module\n",
        "    #     assert (module is not None), \"pass a module\"\n",
        "\n",
        "    #     for name, sub_module in fp16_model.named_modules():\n",
        "    #         if isinstance(sub_module, nn.Linear):\n",
        "    #             lshape = module.weight.shape\n",
        "    #             setattr(module, name, sub_linear(*lshape, **kargs))\n",
        "\n",
        "\n",
        "    def map(self,\n",
        "            name_module=None,\n",
        "            module=None,\n",
        "            share_weights = True,\n",
        "            **kargs):\n",
        "        \"\"\"\n",
        "        Recursively replaces layers in a module with a custom module.\n",
        "\n",
        "        Args:\n",
        "            module (nn.Module): The input module to be modified.\n",
        "            name_module (list, optional): List of names to filter modules.\n",
        "            share_weights (bool, optional): If True, apply the new module on parallel model with shared weights between the original and the adapted model modules.\n",
        "            load_by_layer (bool, optional): If True, loads state_dict layer by layer (used with share_weights).\n",
        "            **kwargs: Additional keyword arguments.\n",
        "                - Modular (callable): The custom module to replace layers.\n",
        "                - type_Mod (type): The type of module to replace.\n",
        "                - config (dict): Configuration for the custom module.\n",
        "\n",
        "        Returns:\n",
        "            nn.Module: The modified module.\n",
        "        \"\"\"\n",
        "        # config:\n",
        "        self.Modular = kargs.get(\"Modular\", self.Modular)\n",
        "        self.type_Mod = kargs.get(\"type_Mod\", self.type_Mod)\n",
        "        config = kargs.get(\"config\", self.config)\n",
        "        config = {} if config is None else config\n",
        "\n",
        "        module = self.module if (module is None) else module\n",
        "        assert (module is not None), \"pass a module\"\n",
        "\n",
        "        if share_weights:\n",
        "          names, sub_modules =  self.Modular_rec(module, name_module, config)\n",
        "          return module, (names, sub_modules)\n",
        "        else:\n",
        "          module_adapt=copy.deepcopy(module)\n",
        "          names, sub_modules = self.Modular_rec(module_adapt , name_module, config)\n",
        "          return module_adapt, (names, sub_modules)\n",
        "\n",
        "\n",
        "    def Modular_rec(self, model, patterns, config):\n",
        "      def replace_module_by_path(model, module_path, new_module):\n",
        "\n",
        "        # Split the module path string\n",
        "        parts = module_path.split('.')\n",
        "\n",
        "        # print(parts)\n",
        "        # Access the module using getattr() with the split parts\n",
        "        current_module = model\n",
        "        for part in parts[:-1]:\n",
        "            current_module = getattr(current_module, part)\n",
        "\n",
        "        # Replace the existing module with the new module\n",
        "        setattr(current_module, parts[-1], new_module)\n",
        "\n",
        "\n",
        "      modules = dict(model.named_modules())\n",
        "      matched_modules = []\n",
        "      matched_pttn = []\n",
        "\n",
        "      def main(model, name):\n",
        "          Modular_module = self.Modular(module, **config)\n",
        "\n",
        "          matched_pttn.append(name), matched_modules.append(Modular_module)\n",
        "\n",
        "          replace_module_by_path(model,name , Modular_module)\n",
        "\n",
        "\n",
        "      if patterns is None:\n",
        "          for name, module in modules.items():\n",
        "            if isinstance(module, self.type_Mod):\n",
        "              main(model, name)\n",
        "\n",
        "      elif isinstance(patterns, list):\n",
        "        for pattern in patterns:\n",
        "            for name, module in modules.items():\n",
        "              if re.search(pattern, name) and isinstance(module, self.type_Mod):\n",
        "                if not name in matched_pttn:\n",
        "                  main(model, name)\n",
        "      elif isinstance(patterns, str):\n",
        "        pattern = patterns\n",
        "        for name, module in modules.items():\n",
        "          if re.search(pattern, name) and isinstance(module, self.type_Mod):\n",
        "            if not name in matched_pttn:\n",
        "              main(model, name)\n",
        "\n",
        "\n",
        "      return matched_pttn, matched_modules\n",
        "\n",
        "\n",
        "    # def Modular_rec(self, module, name_module=None, name=None, parent_name=None, full_names=None):\n",
        "    #     if full_names is None:\n",
        "    #       full_names = []\n",
        "\n",
        "    #     if isinstance(module, self.type_Mod):\n",
        "    #         if name_module is None or (name and name_module is not None and name in name_module):\n",
        "    #             Modular_module = self.Modular(module, self.config)\n",
        "    #             full_names.append(parent_name)\n",
        "    #             return Modular_module, full_names\n",
        "\n",
        "    #     for name, sub_module in module.named_children():\n",
        "\n",
        "    #         # Construct the full name of the submodule\n",
        "    #         full_name = f\"{parent_name}.{name}\" if parent_name else name\n",
        "\n",
        "    #         if isinstance(sub_module, nn.ModuleDict):\n",
        "    #             # Handle nn.ModuleDict by iterating through values\n",
        "    #             for key_sub, value_sub in sub_module.items():\n",
        "    #                 sub_module[key_sub], full_names = self.Modular_rec(value_sub, name_module, name, full_name, full_names)\n",
        "    #         elif isinstance(sub_module, nn.ModuleList):\n",
        "    #             # Handle nn.ModuleList by iterating through elements\n",
        "    #             for i, elem_sub in enumerate(sub_module):\n",
        "    #                 sub_module[i], full_names = self.Modular_rec(elem_sub, name_module, name, full_name, full_names)\n",
        "    #         else:\n",
        "    #             # Handle other cases\n",
        "    #             sub_module, full_names = self.Modular_rec(sub_module, name_module, name, full_name, full_names)\n",
        "\n",
        "\n",
        "    #         setattr(module, name, sub_module)\n",
        "\n",
        "    #     return module, full_names\n",
        "\n",
        "    @staticmethod\n",
        "    def load_state_dict_layer_by_layer(state_dict, module):\n",
        "        \"\"\"\n",
        "        Loop over layers to load weights to a module.\n",
        "\n",
        "        Args:\n",
        "            module (nn.Module): The input module to load state_dict.\n",
        "            state_dict (dict): Dict of weights to load layer by layer in modules.\n",
        "        \"\"\"\n",
        "\n",
        "        for name, param in module.named_parameters():\n",
        "            if name in state_dict:\n",
        "                param.data.copy_(state_dict[name])\n",
        "            else:\n",
        "                print(f\"Warning: Parameter {name} not found in the state_dict.\")\n",
        "\n",
        "\n",
        "    def export_mod(self,\n",
        "                           join=False,\n",
        "                           module= None,\n",
        "                           names=[],\n",
        "                           **kargs):\n",
        "        module = self.module if (module is None) else module\n",
        "        assert (module is not None), \"pass a module\"\n",
        "\n",
        "        \"\"\"\n",
        "        Export the weights of the Modular layers in the module.\n",
        "\n",
        "        Args:\n",
        "            module (nn.Module): The input module to export the weights of Modular state_dict.\n",
        "            join (bool): True if is a OrderedDict or False to a dict of OrderedDicts.\n",
        "        \"\"\"\n",
        "\n",
        "        Modular = kargs.get(\"Modular\", self.Modular)\n",
        "\n",
        "        def concatenate_keys(input_dict, separator='.'):\n",
        "            result_dict = OrderedDict()\n",
        "\n",
        "            for key, value in input_dict.items():\n",
        "                if isinstance(value, dict):\n",
        "                    nested_dict = concatenate_keys(value, separator)\n",
        "                    for nested_key, nested_value in nested_dict.items():\n",
        "                        result_key = f\"{key}{separator}{nested_key}\"\n",
        "                        result_dict[result_key] = nested_value\n",
        "                else:\n",
        "                    result_dict[key] = value\n",
        "\n",
        "            return result_dict\n",
        "\n",
        "\n",
        "        all_weights = {}\n",
        "        for name, module in module.named_modules():\n",
        "            if isinstance(module, Modular):\n",
        "                if len(names)>0:\n",
        "                  print(name)\n",
        "                  if name.split(\".\")[-1] in names:\n",
        "                    all_weights[name] = module.state_dict()\n",
        "                else:\n",
        "                    all_weights[name] = module.state_dict()\n",
        "\n",
        "        if join:\n",
        "          all_weights = concatenate_keys(all_weights)\n",
        "        return all_weights\n",
        "\n",
        "\n",
        "    def load_mod(self,\n",
        "                 all_weights=None,\n",
        "                 module=None,\n",
        "                 **kargs):\n",
        "      \"\"\"\n",
        "      Export the weights of the Modular layers in the module.\n",
        "\n",
        "      Args:\n",
        "          module (nn.Module): The input module to export the weights of Modular state_dict.\n",
        "          all_weights (dict): Dict of  OrderedDicts.\n",
        "\n",
        "      Exemple:\n",
        "          module\n",
        "\n",
        "          module.load_state_dict(all_weights) # if all_weights is OrderedDict to apply.\n",
        "          load_mod_weights(module, all_weights) # if all_weights is a dict of OrderedDict weights.\n",
        "      \"\"\"\n",
        "      module = self.module if (module is None) else module\n",
        "      assert (module is not None), \"pass a module\"\n",
        "\n",
        "\n",
        "      Modular = kargs.get(\"Modular\", self.Modular)\n",
        "      for name, module in module.named_modules():\n",
        "          if isinstance(module, Modular):\n",
        "              module.load_state_dict(all_weights[name])\n",
        "\n",
        "\n",
        "    # target_path be your own\n",
        "    def freeze(self, target_paths=None, freeze=True, module=None):\n",
        "        \"\"\"\n",
        "        Freeze or unfreeze layers in a PyTorch module based on the specified paths.\n",
        "\n",
        "        Args:\n",
        "            module (nn.Module): The PyTorch module.\n",
        "            target_paths (list, optional): A list of paths to the target layers in the module. If None, freeze or unfreeze all layers. Default is None.\n",
        "\n",
        "        Returns:\n",
        "            None\n",
        "        \"\"\"\n",
        "        module = self.module if (module is None) else module\n",
        "        assert (module is not None), \"pass a module\"\n",
        "\n",
        "        freezed = ()\n",
        "        if target_paths is None:\n",
        "            # Freeze or unfreeze all layers in the module\n",
        "            for param in module.parameters():\n",
        "                param.requires_grad = not freeze\n",
        "        else:\n",
        "            # Freeze or unfreeze the specified layers\n",
        "            for target_path in target_paths:\n",
        "                current_module = module\n",
        "                keys = target_path.split('.')\n",
        "\n",
        "                for key in keys[:-1]:\n",
        "                    current_module = getattr(current_module, key)\n",
        "\n",
        "                target_layer = getattr(current_module, keys[-1])\n",
        "\n",
        "                for param in target_layer.parameters():\n",
        "                    param.requires_grad = not freeze\n",
        "\n",
        "    @staticmethod\n",
        "    def meta_fun(objects, methods, args):\n",
        "        for obj, method_name, arg in zip(objects, methods, args):\n",
        "            method = getattr(obj, method_name)\n",
        "            if callable(method):\n",
        "                return method(arg)\n",
        "            else:\n",
        "                print(f\"Object {obj} does not have method {method_name}\")\n",
        "\n",
        "\n",
        "    def activate(self, layer_names=None, key=None, module=None, config={}, **karg):\n",
        "        # print(karg)\n",
        "        module = self.module if (module is None) else module\n",
        "        assert (module is not None), \"pass a module\"\n",
        "\n",
        "        Modular = karg.get('Modular', self.Modular)\n",
        "\n",
        "        returns = ()\n",
        "        for name, module in module.named_modules():\n",
        "            if any(layer_name in name for layer_name in layer_names) and isinstance(module, Modular):\n",
        "              print(f'[{name}] {key} activated')\n",
        "              returns += ((name , getattr(module, key)(**config)), )\n",
        "\n",
        "        return returns\n",
        "\n",
        "    @staticmethod\n",
        "    def get_attribute(obj, query):\n",
        "        keys = query.split('.')\n",
        "        current_obj = obj\n",
        "        for key in keys:\n",
        "            try:\n",
        "                key = int(key)\n",
        "            except ValueError:\n",
        "                pass\n",
        "            if isinstance(current_obj, dict):\n",
        "                current_obj = current_obj.get(key)\n",
        "            elif isinstance(current_obj, list):\n",
        "                if ':' in key:\n",
        "                  slice_obj = slice(*map(int, key.split(':')))\n",
        "                else:\n",
        "                  slice_obj = int(key)\n",
        "\n",
        "                current_obj = current_obj[slice_obj]\n",
        "\n",
        "            else:\n",
        "                # print(current_obj, key)\n",
        "                current_obj = getattr(current_obj, key)\n",
        "        return current_obj\n",
        "\n",
        "\n",
        "    def get(self, layer_names=None, key=None, module=None, **karg):\n",
        "        module = self.module if (module is None) else module\n",
        "        assert (module is not None), \"pass a module\"\n",
        "\n",
        "        Modular = karg.get('Modular', self.Modular)\n",
        "\n",
        "        returns = ()\n",
        "        for name, module in module.named_modules():\n",
        "            if any(layer_name in name for layer_name in layer_names) and isinstance(module, Modular):\n",
        "              att = self.get_attribute(module, key)\n",
        "\n",
        "              returns += ((name , att), )\n",
        "\n",
        "        return returns\n",
        "\n",
        "\n",
        "    def trainable_params(self, non_zero=False, module=None):\n",
        "        \"\"\"\n",
        "        Function to print the parameters of the model and whether they are set for training.\n",
        "\n",
        "        Args:\n",
        "        - module: The PyTorch module.\n",
        "        - non_zero: If True, counts only non-zero parameters.\n",
        "\n",
        "        Returns:\n",
        "        - count: Total number of trainable parameters.\n",
        "        - total: Total number of parameters.\n",
        "        \"\"\"\n",
        "        module = self.module if (module is None) else module\n",
        "        assert module is not None, \"Please pass a module\"\n",
        "\n",
        "        print(f\"{'####'*10}\")\n",
        "        print(f\"{'####'*2}{' Parameters of the Model: ':^5}{'######'}\")\n",
        "        print(f\"{'####'*10}\")\n",
        "        count = 0\n",
        "        total = 0\n",
        "\n",
        "        print(module)\n",
        "        for name, param in module.named_parameters():\n",
        "            if non_zero:\n",
        "                temp = torch.count_nonzero(param).item()\n",
        "            else:\n",
        "                temp = param.numel()\n",
        "\n",
        "            total += temp\n",
        "            if param.requires_grad:\n",
        "              count += temp\n",
        "              print(f\"{name}: Trainable\")\n",
        "            # else:\n",
        "              # print(f\"{name}: No\")\n",
        "\n",
        "        print()\n",
        "        print(f\"Total not trainable: {total-count}\")\n",
        "        print(f\"Trainable parameters {count} ({count / total:.2%}) {'non zero ' if non_zero else ' '}\")\n",
        "        print(f\"{'####'*10}\")\n",
        "\n",
        "        return count, total\n",
        "\n",
        "\n",
        "    def copy_arch(self, module=None):\n",
        "      module = self.module if (module is None) else module\n",
        "      assert (module is not None), \"pass a module\"\n",
        "\n",
        "      def copy_with_shared_weights(model1):\n",
        "          model2 = copy.deepcopy(model1)\n",
        "          for name, param in model1.named_parameters():\n",
        "              model2.get_parameter(name).data = param.data\n",
        "          return model2\n",
        "\n",
        "      return copy_with_shared_weights(module)\n",
        "\n",
        "\n",
        "#######\n",
        "\n",
        "\n",
        "class ModularUtils_class:\n",
        "    def __init__(self, **kargs):\n",
        "        # print(test)\n",
        "        # Looping through named modules and assigning lambda functions\n",
        "        for name, module in self.named_modules():\n",
        "            # Creating a lambda function that calls custom with the module's name\n",
        "            ModularUtils(module=module,\n",
        "                                       tag=name, **kargs)\n",
        "            # print(module.util.module)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "OPW4YHDtncQB",
        "cellView": "code"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## ⛱ Adapter"
      ],
      "metadata": {
        "id": "W1W0sGoQCL-j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def create_sub_dict(original_dict, keys_to_keep):\n",
        "    \"\"\"\n",
        "    Create a sub-dictionary from the original dictionary by keeping only the specified keys.\n",
        "\n",
        "    Args:\n",
        "    - original_dict (dict): The original dictionary.\n",
        "    - keys_to_keep (list): A list of keys to keep in the sub-dictionary.\n",
        "\n",
        "    Returns:\n",
        "    - sub_dict (dict): The sub-dictionary containing only the specified keys.\n",
        "    \"\"\"\n",
        "    sub_dict = {key: original_dict[key] for key in keys_to_keep if key in original_dict}\n",
        "    return sub_dict\n",
        "\n",
        "\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch\n",
        "import copy\n",
        "from collections import OrderedDict\n",
        "\n",
        "\n",
        "class Lora(nn.Module):\n",
        "    def __init__(self, base_linear,\n",
        "                 alpha= 0.5, r = 1,\n",
        "                 **config):\n",
        "        super(Lora, self).__init__()\n",
        "\n",
        "        self.merged = False\n",
        "        self.config = {\n",
        "            'alpha': alpha,\n",
        "            'r': r,\n",
        "            **config  # Merge additional configuration options\n",
        "        }\n",
        "        # dont save as a dict\n",
        "        self.base_linear = base_linear\n",
        "\n",
        "        self.const = nn.Parameter(torch.Tensor([self.config.get(\"alpha\")/self.config.get(\"r\")]), requires_grad=False)\n",
        "        self.config['o'], self.config['i'] = base_linear.weight.shape\n",
        "\n",
        "        self.init()\n",
        "\n",
        "    def init(self):\n",
        "        w = self.base_linear.weight.detach().T\n",
        "\n",
        "        w1 = nn.Linear(self.config[\"i\"],\n",
        "                       self.config[\"r\"], bias=False, dtype=w.dtype)\n",
        "        w1=w1.weight.T\n",
        "\n",
        "        self.loraA = nn.Parameter(w1)\n",
        "\n",
        "\n",
        "        self.loraB = nn.Parameter(torch.zeros(\n",
        "            (self.config[\"r\"], self.config[\"o\"]), dtype=w.dtype)\n",
        "        )\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        return self.base_linear(x) + self.main(x)\n",
        "\n",
        "\n",
        "    def main(self,x):\n",
        "        low = ( (x @ self.loraA) @self.loraB )\n",
        "\n",
        "        return self.const * low\n",
        "\n",
        "\n",
        "\n",
        "#####################\n",
        "    @torch.no_grad()\n",
        "    def merge(self,name= None, arg=None):\n",
        "        if name == None:\n",
        "          name = '.'\n",
        "        if self.merged == True:\n",
        "          return print(f'[{name}]>> already Done')\n",
        "        else:\n",
        "          print(f'[{name}]>> Done!')\n",
        "\n",
        "        dtype = self.loraB.dtype\n",
        "\n",
        "        self.merged = True\n",
        "        base_linear = self.base_linear\n",
        "\n",
        "        base_linear.weight.data = base_linear.weight.to(dtype=dtype)\n",
        "\n",
        "        base_linear.weight[:] += (self.const*self.loraA @ self.loraB).T\n",
        "\n",
        "\n",
        "        # base_linear.weight[:] = (self.m/self.norm()) * base_linear.weight\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    def export(self):\n",
        "        keys_to_keep = [\n",
        "            # 'const',\n",
        "                        'loraA',\n",
        "                        'loraB',\n",
        "                        ]\n",
        "        return create_sub_dict(dict(self.state_dict()),\n",
        "                               keys_to_keep)\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def unmerge(self, arg=None):\n",
        "        if self.merged == True:\n",
        "          return print(f'[{name}]>> already Done')\n",
        "        else:\n",
        "          print(f'[{name}]>> Done!')\n",
        "\n",
        "        dtype = self.loraB.dtype\n",
        "\n",
        "        self.merged = False\n",
        "        base_linear = self.base_linear\n",
        "\n",
        "        base_linear.weight.data = base_linear.weight.to(dtype=dtype)\n",
        "\n",
        "        base_linear.weight[:] -= (self.const*self.loraA @ self.loraB).T\n",
        "\n",
        "        # base_linear.weight[:] = (self.m/self.norm()) * base_linear.weight\n",
        "\n",
        "\n",
        "# class MLora(nn.Module):\n",
        "#   def __init__(self, base_linear, config):\n",
        "#     super(MLora, self).__init__()\n",
        "\n",
        "#   def forward(self, x):\n",
        "#     return x\n",
        "\n"
      ],
      "metadata": {
        "id": "IXzbRfh4CQT4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 🕸 sparcity\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "06akBZ3zC1S8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "https://pytorch.org/tutorials/prototype/semi_structured_sparse.html\n",
        "\n",
        "https://github.com/locuslab/wanda?tab=readme-ov-file\n",
        "\n",
        "https://github.com/VainF/Torch-Pruning?tab=readme-ov-file\n",
        "\n",
        "https://eric-mingjie.github.io/wanda/home.html"
      ],
      "metadata": {
        "id": "7KaUijcADRNE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from transformers import AutoTokenizer, AutoModel\n",
        "# model_name = 'sentence-transformers/paraphrase-MiniLM-L3-v2'\n",
        "# model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "# modelref_s = AutoModel.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "from einops import einsum,pack\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch\n",
        "import copy\n",
        "from collections import OrderedDict\n",
        "\n",
        "\n",
        "\n",
        "def norm_d(tensor, p=2, dim=0, sum_old=None):\n",
        "    sum_tensor = tensor.pow(p).sum(dim=dim)\n",
        "\n",
        "    # If sum_old is provided, accumulate the sum\n",
        "    if sum_old is not None:\n",
        "        sum_tensor += sum_old\n",
        "\n",
        "    return sum_tensor\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# # W: weight matrix (C_out, C_in);\n",
        "# # X: input matrix (N * L, C_in);\n",
        "# # s: desired sparsity, between 0 and 1;\n",
        "\n",
        "# # wanda\n",
        "\n",
        "# def prune(W, X, s):\n",
        "#     metric = W.abs() * X.norm(p=2, dim=0)   # get the Wanda pruning metric\n",
        "\n",
        "#     _, sorted_idx = torch.sort(metric, dim=1)   # sort the weights per output\n",
        "#     pruned_idx = sorted_idx[:,:int(C_in * s)]   # get the indices of the weights to be pruned\n",
        "#     W.scatter_(dim=1, index=pruned_idx, src=0)   # get the index of pruned weights\n",
        "\n",
        "#     return W\n",
        "\n",
        "\n",
        "class Prune(nn.Module):\n",
        "    def __init__(self, base_linear,\n",
        "                s=0.8,acum=False,copyw=False,\n",
        "                 **config):\n",
        "        super(Prune, self).__init__()\n",
        "        self.config = {\n",
        "                's': s,\n",
        "                'acum':acum,\n",
        "                'copyw':copyw,\n",
        "                **config\n",
        "            }\n",
        "\n",
        "        if self.config['copyw']:\n",
        "          self.main = copy.deepcopy(base_linear)\n",
        "        else:\n",
        "          self.main = base_linear\n",
        "        self.quantized = False\n",
        "\n",
        "\n",
        "        self.sum_old = 0\n",
        "\n",
        "        self.mask = nn.Parameter(torch.tensor([1], dtype=torch.bool),\n",
        "                                 requires_grad=False)\n",
        "\n",
        "        self.pruned = False\n",
        "\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def forward(self, x):\n",
        "        W = self.main.weight.detach()\n",
        "        y = self.main(x)\n",
        "        y_hat = self.filter_prune(W,x)\n",
        "        # y_hat = self.filter_prune(W,x)\n",
        "\n",
        "        # print(self.sum_old)\n",
        "        print(torch.mean((y - y_hat)**2))\n",
        "        return y_hat\n",
        "\n",
        "    # def fake_prune(self, X):\n",
        "    #     W = self.main.weight.detach()\n",
        "    #     return  einsum((W*self.mask), X, \"out inn, b l inn -> b l out\")\n",
        "\n",
        "\n",
        "    def prune(self):\n",
        "      W = self.main.weight.detach()\n",
        "\n",
        "      if self.pruned:\n",
        "        print(\"already pruned!\")\n",
        "      else:\n",
        "        self.pruned = True\n",
        "        mask = self.mask\n",
        "\n",
        "        if mask.dtype is not torch.bool:\n",
        "          mask = mask.to(torch.bool)\n",
        "\n",
        "        W[~mask] = 0\n",
        "\n",
        "      # self.main = self.main\n",
        "\n",
        "    def filter_prune(self, W, X):\n",
        "        s = self.config['s']\n",
        "        W = self.main.weight.detach()\n",
        "\n",
        "        C_in = W.shape[1]\n",
        "        X_join = pack(X, '* d')\n",
        "\n",
        "        if self.config.get('acum'):\n",
        "          self.sum_old = norm_d(X_join[0],\n",
        "                                p=2, dim= 0,\n",
        "                                sum_old= self.sum_old)\n",
        "          norm_vec = self.sum_old.pow(1/2)\n",
        "        else:\n",
        "          norm_vec = X_join[0].norm(p=2, dim=0)\n",
        "\n",
        "\n",
        "        metric = W.abs() * norm_vec    # get the Wanda pruning metric\n",
        "        _, sorted_idx = torch.sort(metric, dim=1)   # sort the weights per output\n",
        "        pruned_idx = sorted_idx[:,:int(C_in * s)]   # get the indices of the weights to be pruned\n",
        "        self.pruned_idx = pruned_idx\n",
        "\n",
        "        mask = torch.ones_like(W, dtype=W.dtype)\n",
        "\n",
        "        mask.scatter_(dim=1, index=self.pruned_idx,\n",
        "                      src=torch.zeros_like(self.pruned_idx, dtype=W.dtype))\n",
        "\n",
        "\n",
        "        self.mask.data = mask\n",
        "\n",
        "        return X @ (W*mask).T\n",
        "\n",
        "\n",
        "# Mod_prune = ModularUtils(Modular=Prune,\n",
        "#                          config_Mod= {\n",
        "#                              's': 0.8,\n",
        "#                              'acum':True\n",
        "#                          })\n",
        "\n",
        "# model_base = copy.deepcopy(modelref_s)\n",
        "# model_mod = Mod_prune.copy_arch(model_base)\n",
        "\n",
        "\n",
        "# base_w = model_base.get_submodule('encoder.layer')[-1].output.dense\n",
        "# base_w, base_w.weight\n",
        "\n",
        "# a = Mod_prune.app_Modular(\n",
        "#     model_mod.get_submodule('encoder.layer')[-1].output,\n",
        "#     ['dense']\n",
        "# )\n",
        "\n",
        "# mod_w = model_mod.get_submodule('encoder.layer')[-1].output.dense\n",
        "# mod_w, mod_w.base_linear.weight\n",
        "\n",
        "# model_mod(**tokenize_batch_s([\n",
        "#     'hi',\n",
        "#     'hi',\n",
        "#     'hi',\n",
        "#     'hi2'\n",
        "#                                ]))\n",
        "\n",
        "# # Mod_prune.activate_layers(\n",
        "# #     module = model_mod,\n",
        "# #     layer_names = ['dense'],\n",
        "# #     activation_func='prune')\n",
        "\n",
        "# # a[0].dense.pruned_idx\n",
        "\n",
        "# # a[0].dense.prune()\n",
        "\n"
      ],
      "metadata": {
        "id": "h2JDBwJAWh_p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### prune + Lora"
      ],
      "metadata": {
        "id": "vb7pixXhID0c"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # faz um foward e progapaga o input com a mascara para guarda em todas as\n",
        "# # layers com Prune, os indexs para transformar os valores em zero.\n",
        "\n",
        "\n",
        "# class ModMix(nn.Module):\n",
        "#     def __init__(self, base_linear, config):\n",
        "#         super(ModMix, self).__init__()\n",
        "#         self.Mix = nn.ModuleDict({\n",
        "#             'base_linear': Prune(base_linear, config)\n",
        "#         })\n",
        "\n",
        "#         self.name = 'adapter_a'\n",
        "#         self.Mix[self.name] = Lora(base_linear, config)\n",
        "#         # self.Mix\n",
        "\n",
        "#     def forward(self, x):\n",
        "#       return self.Mix[self.name](x)\n",
        "\n",
        "#     def set(self, name):\n",
        "#       self.name = name\n"
      ],
      "metadata": {
        "id": "5XOygUTyIDIt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 🔰 quantization"
      ],
      "metadata": {
        "id": "MWWUVgJFApw9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://pytorch.org/tutorials/intermediate/dynamic_quantization_bert_tutorial.html\n",
        "\n",
        "https://pytorch.org/tutorials/recipes/quantization.html\n",
        "\n",
        "https://pytorch.org/docs/stable/quantization.html\n",
        "\n",
        "> https://github.com/SqueezeAILab/SqueezeLLM#from-scratch-quantization"
      ],
      "metadata": {
        "id": "58bBdoGYBuTU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import onnx\n",
        "# from onnxruntime.quantization import quantize_dynamic, QuantType\n",
        "\n",
        "# model_fp32 = 'path/to/the/model.onnx'\n",
        "# model_quant = 'path/to/the/model.quant.onnx'\n",
        "# quantized_model = quantize_dynamic(model_fp32, model_quant)\n"
      ],
      "metadata": {
        "id": "E8a41a5bsOFw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# from transformers import AutoTokenizer, AutoModel\n",
        "\n",
        "# model_name = 'sentence-transformers/paraphrase-MiniLM-L3-v2'\n",
        "# model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "# modelref_s = AutoModel.from_pretrained(model_name)\n",
        "\n",
        "import copy\n",
        "def quantize_tensor(x, qmin, qmax, bits, symmetric=True, z=None):\n",
        "\n",
        "    bits_n = 2 ** bits - 1\n",
        "    s = (qmax - qmin) / bits_n\n",
        "\n",
        "    if symmetric:\n",
        "        if z is None:\n",
        "            z = torch.round(torch.mean(x))\n",
        "        xc = x\n",
        "        q = int(bits_n / 2)\n",
        "        qx = torch.clamp(torch.round(xc / s) - z, -q, q + 1)\n",
        "    else:\n",
        "        xc = (x - qmin) / (qmax - qmin)\n",
        "        z = 0 if z is None else z\n",
        "        qx = torch.clamp(torch.round(xc / s) - z, 0, bits_n)\n",
        "\n",
        "    return qx, s\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def fake_quant(W, dim=None, mask=None, bits=8, symmetric=True):\n",
        "  z = 0\n",
        "\n",
        "  if dim is None:\n",
        "    qmin = torch.min(W)\n",
        "    qmax = torch.max(W)\n",
        "  else:\n",
        "    qmin = torch.min(W, dim, keepdims=True).values\n",
        "    qmax = torch.max(W, dim, keepdims=True).values\n",
        "\n",
        "  qW, s = quantize_tensor(W,\n",
        "                            qmin, qmax,\n",
        "                            bits=bits,\n",
        "                            symmetric= symmetric,\n",
        "                            z = z)\n",
        "\n",
        "  print(torch.mean((W - (qW * s))**2))\n",
        "  print('done!')\n",
        "  return qW, (s, qmin, qmax)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#############################\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class QLinear(nn.Module):\n",
        "    def __init__(self, qmin, qmax,\n",
        "                 bits, s, W, z=None, mask=None):\n",
        "        super(QLinear, self).__init__()\n",
        "        self.bits = bits\n",
        "        self.qmin = nn.Parameter(qmin, requires_grad=False)\n",
        "        self.qmax = nn.Parameter(qmax, requires_grad=False)\n",
        "        self.s = nn.Parameter(s, requires_grad=False)\n",
        "        self.z = z\n",
        "        self.W = W\n",
        "\n",
        "    def forward(self, x):\n",
        "      return self.W(x)\n",
        "\n",
        "    def __repr__(self):\n",
        "        qshape = list(self.qmin.data.shape)\n",
        "        sshape = list(self.s.data.shape)\n",
        "        if len(sshape)==0:\n",
        "          qshape = sshape = [1]\n",
        "\n",
        "        if (self.z is not None):\n",
        "          zshape = list(self.z.data.shape)\n",
        "          return f'''QLinear(W={self.W},\\n\\tbits={self.bits}, qmin_max={qshape}, z={zshape}, s={sshape})'''\n",
        "        else:\n",
        "          return f'''QLinear(W={self.W},\\n\\tbits={self.bits}, qmin_max={qshape}, s={sshape})'''\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def export(self, **args):\n",
        "      # self.W.weight.detach()\n",
        "      self.W.weight.requires_grad = False\n",
        "\n",
        "      if self.W.bias is not None:\n",
        "        self.W.bias.data = self.W.bias.to(torch.float16)\n",
        "\n",
        "      self.W.weight.requires_grad = False\n",
        "\n",
        "      self.W.weight.data[:] = self.W.weight.data / self.s.data\n",
        "      w = self.W.weight.detach().to(dtype=torch.int8)\n",
        "\n",
        "      self.W.weight.data = w\n",
        "\n",
        "      self.s.data = self.s.to(dtype=torch.float16)\n",
        "      self.qmin.data = self.qmin.to(dtype=torch.float16)\n",
        "      self.qmax.data = self.qmax.to(dtype=torch.float16)\n",
        "\n",
        "\n",
        "##################################\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class Quantize_Linear(nn.Module):\n",
        "    def __init__(self, base_linear,\n",
        "                 bits = 8, acum= False, symmetric= True, copyw=False,\n",
        "                 **config):\n",
        "        super(Quantize_Linear, self).__init__()\n",
        "\n",
        "        self.config = {\n",
        "                'bits': bits,\n",
        "                'acum':acum,\n",
        "                \"symmetric\": symmetric,\n",
        "                \"copyw\":copyw,\n",
        "                **config\n",
        "            }\n",
        "        # print(self.config)\n",
        "\n",
        "        # print(config)\n",
        "\n",
        "        self.merged = False\n",
        "        if self.config['copyw']:\n",
        "          self.main = copy.deepcopy(base_linear)\n",
        "        else:\n",
        "          self.main = base_linear\n",
        "        self.quantized = False\n",
        "\n",
        "    def forward(self, x):\n",
        "      return self.main(x)\n",
        "\n",
        "\n",
        "    def quantize(self, dim=None, mask=None, **kargs):\n",
        "        # kargs\n",
        "        # print(self.config.get('bits'))\n",
        "\n",
        "        if self.quantized:\n",
        "          print(\"already quantized!\")\n",
        "        else:\n",
        "          self.quantized = True\n",
        "          linear = self.main\n",
        "          bits = self.config.get('bits')\n",
        "          symmetric = self.config.get('symmetric')\n",
        "          W=linear.weight.detach()\n",
        "          dim = 1 if dim is None else dim\n",
        "\n",
        "          if mask is not None:\n",
        "\n",
        "            ds = W.shape[dim]\n",
        "\n",
        "            l = [1]*ds\n",
        "            qmins = []\n",
        "            qmaxs = []\n",
        "            ss = []\n",
        "\n",
        "            wm = zip(W.split(l,dim), mask.split(l,dim))\n",
        "\n",
        "\n",
        "            for ww, m in wm:\n",
        "              ws = ww[m]\n",
        "              if len(ws)>1:\n",
        "                wq, (s, qmin, qmax) = fake_quant(ws,0, bits=bits, symmetric=symmetric)\n",
        "              else:\n",
        "                if len(ws)==0:\n",
        "                  ws = torch.tensor([0])\n",
        "                  s = torch.tensor([1])\n",
        "                  qmin = qmax = ws\n",
        "                else:\n",
        "                  qmin = s = qmax = ws\n",
        "                wq = 1\n",
        "\n",
        "              ww[m] = wq *s\n",
        "              ww[~m] = 0\n",
        "\n",
        "              qmins.append(qmin)\n",
        "              qmaxs.append(qmax)\n",
        "              ss.append(s)\n",
        "\n",
        "            qmin = torch.cat(qmins)\n",
        "            qmax = torch.cat(qmaxs)\n",
        "            s = torch.cat(ss)\n",
        "\n",
        "            layer = QLinear(qmin = qmin, qmax = qmax,\n",
        "                     bits = bits, s = s, W=linear, mask=True)\n",
        "\n",
        "          else:\n",
        "            w, (s, qmin, qmax) = fake_quant(W, dim = dim, bits=bits, symmetric=symmetric)\n",
        "            W[:] = w*s\n",
        "            layer = QLinear(qmin = qmin, qmax = qmax,\n",
        "                      bits = bits, s = s, W=linear)\n",
        "\n",
        "          self.main = layer\n",
        "\n"
      ],
      "metadata": {
        "id": "752hhOseIWmf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### quantization + lora"
      ],
      "metadata": {
        "id": "JGUDy8-MHPSb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from einops import einsum\n",
        "# from einops import pack\n",
        "\n",
        "\n",
        "# import torch.nn as nn\n",
        "# import torch\n",
        "# import copy\n",
        "# from collections import OrderedDict\n",
        "\n",
        "\n",
        "# # faz um foward e progapaga o input com a mascara para guarda em todas as\n",
        "# # layers com Prune, os indexs para transformar os valores em zero.\n",
        "\n",
        "\n",
        "# class ModMix(nn.Module):\n",
        "#     def __init__(self, base_linear, config):\n",
        "#         super(ModMix, self).__init__()\n",
        "#         self.Mix = nn.ModuleDict({\n",
        "#             'base_linear': Quantize_Linear(base_linear, config)\n",
        "#         })\n",
        "\n",
        "#         self.name = 'adapter_a'\n",
        "#         self.Mix[self.name] = Lora(self.Mix['base_linear'].qlinear.W, config)\n",
        "#         # self.Mix\n",
        "\n",
        "#     def forward(self, x):\n",
        "#       return self.Mix[self.name](x)\n",
        "\n",
        "#     def set(self, name):\n",
        "#       self.name = name\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# def freeze_model(model):\n",
        "#     for param in model.parameters():\n",
        "#         param.requires_grad = False\n",
        "\n",
        "# freeze_model(modelref_s)\n",
        "\n",
        "\n",
        "# Mod_quant = ModularUtils(Modular=ModMix,\n",
        "#                          config_Mod= {\n",
        "#                              'bits': 3,\n",
        "#                              'alpha':0.2,\n",
        "#                              'r':2,\n",
        "#                              'call':True\n",
        "#                             #  'acum':True\n",
        "#                          })\n",
        "\n",
        "# model_base = copy.deepcopy(modelref_s)\n",
        "# model_mod = Mod_quant.copy_arch(model_base)\n",
        "\n",
        "\n",
        "# sub = model_mod.get_submodule('encoder.layer')[-1]\n",
        "# a = Mod_quant.app_Modular(sub,['output.dense'])\n",
        "\n",
        "# a[0]"
      ],
      "metadata": {
        "id": "6jNn3iG8g6i_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# Some trys\n",
        "\n",
        "# import torch\n",
        "\n",
        "# # Algorithm 1: LoftQ\n",
        "# # input Pre-trained weight W, target rank r, N-bit quantization function qN(.), alternating step T\n",
        "\n",
        "# def LoftQ(W, r, qN, T):\n",
        "#     # Initialize\n",
        "#     A = torch.zeros(W.size(0), r)\n",
        "#     B = torch.zeros(r, W.size(1))\n",
        "\n",
        "#     for t in range(T):\n",
        "#         # Obtain quantized weight\n",
        "#         Qt = qN(W - torch.matmul(A, B))\n",
        "\n",
        "#         # Obtain low-rank approximation\n",
        "#         U, S, V = torch.svd(W - Qt)\n",
        "#         At = U[:, :r]\n",
        "#         Bt = torch.matmul(torch.diag(S[:r]), V[:r, :])\n",
        "\n",
        "#         # Update A and B\n",
        "#         A = At\n",
        "#         B = Bt\n",
        "\n",
        "#     return Qt, At, Bt\n",
        "\n",
        "# # Usage\n",
        "# # Assuming W is your pre-trained weight tensor\n",
        "# # r is your target rank\n",
        "# # qN is your N-bit quantization function\n",
        "# # T is the number of alternating steps\n",
        "\n",
        "# # Qt, At, Bt = LoftQ(W, r, qN, T)\n",
        "\n",
        "\n",
        "# import torch\n",
        "\n",
        "# #unifom\n",
        "# def quantization(x, bits=8):\n",
        "#     # Assuming x is a PyTorch tensor\n",
        "#     # Quantize the input tensor to a specified number of bits\n",
        "\n",
        "#     # Determine the quantization range based on the number of bits\n",
        "#     qnt_max = torch.max(x)\n",
        "#     qnt_min = torch.min(x)\n",
        "#     bits = 2 ** bits - 1\n",
        "\n",
        "#     s = (qnt_max-qnt_max)/bits\n",
        "\n",
        "#     # Scale the input tensor to fit within the quantization range\n",
        "#     x_scaled = torch.clamp(x * s, quant_min, quant_max)\n",
        "\n",
        "#     # Round the scaled tensor to the nearest integer\n",
        "#     x_rounded = torch.round(x_scaled)\n",
        "\n",
        "#     # Scale back to the original range\n",
        "#     x_quantized = x_rounded / quant_max\n",
        "\n",
        "#     return x_quantized\n",
        "\n",
        "# # Example usage\n",
        "# # Assuming x is your input tensor and bits is the number of bits for quantization\n",
        "# # x_quantized = quantization(x, bits)\n",
        "\n",
        "# `\n",
        "\n",
        "# import torch\n",
        "# import torch.nn as nn\n",
        "\n",
        "# import bitsandbytes as bnb\n",
        "# from bnb.nn import Linear4bit\n",
        "\n",
        "# fp16_model = nn.Sequential(\n",
        "#     nn.Linear(64, 64),\n",
        "#     nn.Linear(64, 64)\n",
        "# )\n",
        "\n",
        "# quantized_model = nn.Sequential(\n",
        "#     Linear4bit(64, 64),\n",
        "#     Linear4bit(64, 64)\n",
        "# )\n",
        "\n",
        "# quantized_model.load_state_dict(fp16_model.state_dict())\n",
        "# quantized_model = quantized_model.to(0) # Quantization happens here"
      ],
      "metadata": {
        "id": "6nJiD4Ujko7q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 🗿 MIX"
      ],
      "metadata": {
        "id": "Prgx47fxTC8L"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ModMix(nn.Module):\n",
        "    def __init__(self, base_linear, **config):\n",
        "        super(ModMix, self).__init__()\n",
        "        self.config = config\n",
        "\n",
        "        self.tag = 'base_linear'\n",
        "        self.main_tag = 'Base'\n",
        "        self.name_adapt = 'main'\n",
        "        self.merged = False\n",
        "\n",
        "        self.Mix = nn.ModuleDict({\n",
        "            'base_linear': base_linear\n",
        "        })\n",
        "\n",
        "        self.main = self.Mix['base_linear']\n",
        "\n",
        "        self.set_dict = {\n",
        "            'Prune': lambda base_linear, config: Prune(base_linear,\n",
        "                                               **config),\n",
        "            'Quant': lambda base_linear, config: Quantize_Linear(base_linear,**config),\n",
        "            'Adapt': {\n",
        "                'new': lambda base_linear, config: Lora(base_linear,\n",
        "                                                **config),\n",
        "                'available': dict()\n",
        "                }\n",
        "            }\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "      return self.main(x)\n",
        "\n",
        "    def merge(self):\n",
        "      if self.merged:\n",
        "        return print('Already Done!')\n",
        "\n",
        "      main = []\n",
        "      mask = None\n",
        "      self.mask = mask\n",
        "      self.main_tag = 'base_linear'\n",
        "\n",
        "      self.Mix['base_linear']\n",
        "      names = list(self.Mix.keys() )\n",
        "\n",
        "      for name in names:\n",
        "\n",
        "        if name == 'Prune':\n",
        "          self.mask = copy.deepcopy(self.Mix['Prune'].mask)\n",
        "          mask = self.mask.data.to(torch.bool)\n",
        "          del self.Mix['Prune']\n",
        "\n",
        "        elif name == 'Quant':\n",
        "          self.Mix['Quant'].quantize(1, mask)\n",
        "          self.Mix['base_linear'] = copy.deepcopy(self.Mix['Quant'].main)\n",
        "          del self.Mix['Quant']\n",
        "          del self.mask\n",
        "\n",
        "      self.merged  = True\n",
        "      self.main = self.Mix['base_linear']\n",
        "      self.tag = 'pruned+quantized'\n",
        "\n",
        "\n",
        "#####################################################\n",
        "\n",
        "    def set(self, name, **config):\n",
        "      base = config.get('base', 'base_linear')\n",
        "      if base == \"Prune\":\n",
        "        base_linear = self.Mix[base].main\n",
        "      else:\n",
        "        base_linear = self.Mix[base]\n",
        "      set_value = self.set_dict.get(name, None)\n",
        "\n",
        "      config = self.config if config is None else config\n",
        "      # print(config)\n",
        "\n",
        "\n",
        "      if self.merged:\n",
        "        base_linear = base_linear.W\n",
        "        self.tag = 'pruned+quantized'\n",
        "        self.main_tag = 'base_linear'\n",
        "\n",
        "\n",
        "\n",
        "      if name == \"Prune\":\n",
        "        # print(config)\n",
        "        if config.get(\"fake_prune\", False):\n",
        "          # print(\"config\")\n",
        "          # print(config)\n",
        "          self.main = self.Mix[\"Prune\"].main\n",
        "        else:\n",
        "          self.Mix[\"Prune\"] = set_value(base_linear, config)\n",
        "          self.Mix[\"Prune\"].main  = base_linear\n",
        "          self.main = self.Mix[\"Prune\"]\n",
        "\n",
        "\n",
        "        self.tag+= '+prune'\n",
        "        self.main_tag = 'Prune'\n",
        "\n",
        "\n",
        "      elif name == \"Quant\":\n",
        "        self.Mix[\"Quant\"] = set_value(base_linear, config)\n",
        "        self.main = self.Mix[\"Quant\"]\n",
        "        self.tag+= '+quant'\n",
        "        self.main_tag = 'Quant'\n",
        "\n",
        "      elif name == \"Base\":\n",
        "        self.main = self.Mix[\"base_linear\"]\n",
        "        self.main_tag = 'base_linear'\n",
        "\n",
        "\n",
        "      elif name == \"Adapt\":\n",
        "\n",
        "        self.set_dict['Adapt']['available'][name] = self.set_dict['Adapt']['new'](base_linear, config)\n",
        "\n",
        "        self.set_dict['Adapt']['active'] = self.set_dict['Adapt']['available'][name]\n",
        "        self.Mix[\"Adapt\"] = self.set_dict['Adapt']['active']\n",
        "        self.name_adapt = name\n",
        "\n",
        "        self.set_dict['Adapt']['active'].base_linear = base_linear\n",
        "        self.main = self.set_dict['Adapt']['active']\n",
        "\n",
        "        ver = self.tag.split('+')\n",
        "        if \"adapt\" not in  ver[-1]:\n",
        "          self.tag+= '+adapt'\n",
        "        self.main_tag = 'Adapt'\n",
        "\n",
        "\n",
        "      else:\n",
        "        if (name not in self.set_dict['Adapt']['available'].keys()):\n",
        "          self.set_dict['Adapt']['available'][name] = self.set_dict['Adapt']['new'](base_linear, config)\n",
        "\n",
        "        self.set_dict['Adapt']['active'] = self.set_dict['Adapt']['available'][name]\n",
        "        self.Mix[\"Adapt\"] = self.set_dict['Adapt']['active']\n",
        "        # self.main = self.Mix[\"Adapt\"]\n",
        "        # print(self.set_dict['Adapt']['active'].main)\n",
        "\n",
        "        self.set_dict['Adapt']['active'].base_linear = base_linear\n",
        "        self.main = self.set_dict['Adapt']['active']\n",
        "\n",
        "        self.name_adapt = name\n",
        "\n",
        "        ver = self.tag.split('+')\n",
        "        if f\"adapt[{name}]\" not in  ver[-1]:\n",
        "          self.tag+= f'+adapt[{name}]'\n",
        "        self.main_tag = 'Adapt'\n",
        "\n",
        "\n",
        "    def del_layer(self, name, base):\n",
        "      self.Mix[base] = copy.deepcopy(self.Mix[base])\n",
        "      self.main = self.Mix[base].main\n",
        "\n",
        "      if isinstance(name, list):\n",
        "        for n in name:\n",
        "          del self.Mix[n]\n",
        "      else:\n",
        "        del self.Mix[name]\n",
        "\n",
        "\n",
        "#####################################################\n",
        "\n",
        "    def __repr__(self):\n",
        "      # return self.main.__repr__()\n",
        "        p = self.Mix.__repr__()\n",
        "        p= p.replace('\\n ','\\n\\t  ')\n",
        "        p= p.replace('\\n)','\\n\\t)')\n",
        "        p = re.sub('(Adapt\\w?\\))',f'\\\\1 [{self.name_adapt}]', p)\n",
        "\n",
        "        return f'''Active(main={self.main_tag}\\n\\ttag={self.tag},\\n\\tMix={p}\\n)'''"
      ],
      "metadata": {
        "id": "gXxwXuwZTtS-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## examples"
      ],
      "metadata": {
        "id": "Z0MMG3AJCJIM"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ex1"
      ],
      "metadata": {
        "id": "ZoQE4LrKFo6O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "import copy\n",
        "import re\n",
        "\n",
        "lin = nn.Linear(5,10, True)\n",
        "\n",
        "d = ModMix(lin)\n",
        "\n",
        "# d.set('test')\n",
        "# d(x)\n",
        "\n",
        "# d.main.init()\n",
        "\n",
        "# d.main.merge()\n",
        "\n",
        "# print(d.Mix['base_linear'].weight)\n",
        "\n",
        "# d.set('Quant')\n",
        "# # d.merge()\n",
        "# d(x)\n",
        "\n",
        "# d.main.quantize(0)\n",
        "\n",
        "# d.main.main.W.weight\n",
        "\n",
        "\n",
        "d.set('Prune')\n",
        "x = torch.rand((2,1,5))\n",
        "\n",
        "d(x)\n",
        "d.main.prune()\n",
        "\n",
        "d, d.main.mask\n",
        "\n",
        "d.main.mask.data\n",
        "\n",
        "\n",
        "# print(d.Mix['base_linear'].weight)\n",
        "# d\n",
        "# print(d.main)\n",
        "# # print(d.Mix['base_linear'].W.weight)\n",
        "\n",
        "# d.set('Base')\n",
        "\n",
        "# d"
      ],
      "metadata": {
        "id": "inBDlvVjCLCA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ex2"
      ],
      "metadata": {
        "id": "-_QbTgYyGpZt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class subsub(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(subsub, self).__init__()\n",
        "        self.a = nn.Linear(2, 3)\n",
        "        self.b = nn.Linear(3, 3)\n",
        "        self.sss = nn.Sequential(\n",
        "            nn.Linear(3, 4),\n",
        "            nn.Linear(4, 5),\n",
        "            self.b,\n",
        "        )\n",
        "        self.ff = nn.Linear(5, 2)\n",
        "\n",
        "\n",
        "class ComplexModel(nn.Module,\n",
        "                  ModularUtils_class\n",
        "                   ):\n",
        "    def __init__(self):\n",
        "        super(ComplexModel, self).__init__()\n",
        "\n",
        "\n",
        "        self.layer1 = nn.Linear(2, 3)\n",
        "        self.layer2 = nn.Linear(3, 3)\n",
        "        self.submodule = nn.Sequential(\n",
        "            nn.Linear(3, 4),\n",
        "            nn.Linear(4, 5),\n",
        "            self.layer2,\n",
        "        )\n",
        "        self.deep = nn.ModuleList([\n",
        "            subsub(),\n",
        "            subsub()\n",
        "        ]\n",
        "        )\n",
        "        self.layer3 = nn.Linear(5, 2)\n",
        "\n",
        "        ModularUtils_class.__init__(self,\n",
        "                                    Modular = Lora)  # Manually call enhance modules class constructor\n",
        "\n",
        "#######\n",
        "\n",
        "\n",
        "\n",
        "# Mod = ModularUtils(Modular=Lora)\n",
        "model = ComplexModel()\n",
        "\n",
        "model.layer1.util.module"
      ],
      "metadata": {
        "id": "E5UN3V3fjOff"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Mod = ModularUtils(Modular=Lora)\n",
        "named = lambda n: print(f\"{'#'*50}\\n{'# '*5}{n:^20} #\\n{'#'*50}\")\n",
        "\n",
        "model = ComplexModel()\n",
        "\n",
        "\n",
        "model_base = copy.deepcopy(model)\n",
        "model_mod = model.util.copy_arch(model_base)\n",
        "\n",
        "\n",
        "sub = model_mod.get_submodule(\"deep\")\n",
        "\n",
        "k, n = sub.util.map(\n",
        "    name_module = [\"a\", 'ff'],\n",
        "    config_Mod={'alpha':10}\n",
        ")\n",
        "\n",
        "# k.util.module\n",
        "k.util.trainable_params(True)\n",
        "named('lora_module')\n",
        "print()\n",
        "print(n[0])\n",
        "print(k)\n",
        "\n",
        "\n",
        "named('model_base')\n",
        "print(model_base)\n",
        "\n",
        "named('model_mod')\n",
        "print(model_mod)\n",
        "\n",
        "# # Export weights\n",
        "all_weights = model_mod.util.export_mod( join= False)\n",
        "\n",
        "# # Load weights back\n",
        "model_mod.util.load_mod(all_weights)\n",
        "\n",
        "# all_weights = model_mod.util.export_mod( join= True)\n",
        "# a = model_mod.load_state_dict(all_weights, strict=False)\n",
        "# a.missing_keys\n",
        "\n",
        "\n",
        "\n",
        "ref = k[0].a\n",
        "# ref = k.deep[0].a\n",
        "print(\"#########################\")\n",
        "\n",
        "\n",
        "named(\"k[0].a.base_linear.weight\")\n",
        "print(ref.base_linear.weight)\n",
        "\n",
        "ref.merge()\n",
        "# k[0].a.unmerge()\n",
        "\n",
        "assert torch.allclose(\n",
        "    ref.base_linear.weight,\n",
        "    model_base.deep[0].a.weight) , \"model_base.deep[0].a.weight differ\"\n",
        "\n",
        "assert torch.allclose(\n",
        "    ref.base_linear.weight,\n",
        "    model_mod.deep[0].a.base_linear.weight), \"model_mod.deep[0].a.base_linear.weight differ\"\n",
        "\n",
        "# model_base.deep[0].a.weight.data[:] = 0\n",
        "\n",
        "# ref.base_linear['linear'].weight.data[:] = 0\n",
        "\n",
        "# print(ref.base_linear['linear'].weight)\n",
        "\n",
        "named(\"model_base.deep[0].a.weight\")\n",
        "print(model_base.deep[0].a.weight)"
      ],
      "metadata": {
        "id": "-9llS-VPHEyS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = ComplexModel()\n",
        "\n",
        "model_base = copy.deepcopy(model)\n",
        "model_mod = model.util.copy_arch(model_base)\n",
        "\n",
        "\n",
        "sub = model_mod.get_submodule(\"deep\")\n",
        "\n",
        "k, n = sub.util.map(\n",
        "    name_module = [\"a\", 'ff'], share_weights=False\n",
        ")\n",
        "\n",
        "\n",
        "print(\"######\"*10)\n",
        "print(\"## lora_module ##\")\n",
        "print(n)\n",
        "print(k)\n",
        "\n",
        "print(\"######\"*10)\n",
        "print(\"## model_base ##\")\n",
        "# print(\"\\nModel ori:\")\n",
        "print(model_base)\n",
        "\n",
        "print(\"######\"*10)\n",
        "print(\"## model_mod ##\")\n",
        "# print(\"\\nModel ori:\")\n",
        "print(model_mod)\n",
        "\n",
        "# # Export weights\n",
        "# model.util.module\n",
        "\n",
        "all_weights = k.util.export_mod(join= False)\n",
        "\n",
        "# # Load weights back\n",
        "k.util.load_mod(all_weights, model)\n",
        "\n",
        "\n",
        "ref = k[0].a\n",
        "# ref = k.deep[0].a\n",
        "print(\"#########################\")\n",
        "print(\"k[0].a.base_linear.weight\")\n",
        "print(ref.base_linear.weight)\n",
        "ref.merge()\n",
        "# k[0].a.unmerge()\n",
        "print(\"------- k[0].a.base_linear.weight -------\")\n",
        "print(ref.base_linear.weight)\n",
        "\n",
        "print(\"------- model_base.deep[0].a.weight -------\")\n",
        "print(model_base.deep[0].a.weight)\n"
      ],
      "metadata": {
        "id": "A7Y6pwd1G7qM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### bert example"
      ],
      "metadata": {
        "id": "Kg877UW39n9p"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # Load model directly\n",
        "# from transformers import AutoTokenizer, AutoModelForMaskedLM\n",
        "\n",
        "# tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
        "# model = AutoModelForMaskedLM.from_pretrained(\"bert-base-uncased\")\n",
        "\n",
        "# model"
      ],
      "metadata": {
        "id": "yQKfPwa10Qc4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Mod = ModularUtils()\n",
        "\n",
        "\n",
        "# # model = ComplexModel()\n",
        "\n",
        "# Mod.app_Modular(\n",
        "#     model, # Modular\n",
        "#                 [\"query\", 'key'] # linear layes (type_layers)\n",
        "#                 )\n",
        "\n",
        "\n",
        "# model.get_submodule('bert.encoder')"
      ],
      "metadata": {
        "id": "meRndRiM9m9D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 🧠 Model Arch"
      ],
      "metadata": {
        "id": "QTzDkyn53Nwt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### util functions\n"
      ],
      "metadata": {
        "id": "t74soB3f3QtD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# Example usage:\n",
        "# file_path = \"/content/content/slr_data_tcc/synergy/Smid_2020.csv\"\n",
        "# df = preprocess_csv_file(file_path)\n",
        "\n",
        "\n",
        "# df_train = df.groupby('label').sample(4)\n",
        "# df_temp = df[~df.index.isin(df_train.index)]\n",
        "\n",
        "# df_eval = df_temp.groupby('label').sample(2)\n",
        "# df_test = df_temp[~df_temp.index.isin(df_eval.index)]\n",
        "\n",
        "# df['label'].value_counts()\n",
        "\n",
        "# print(df_train.iloc[0])\n",
        "# df_train.iloc[0]['text']\n",
        "\n",
        "\n",
        "\n",
        "from einops import unpack\n",
        "from itertools import chain\n",
        "from collections import namedtuple\n",
        "\n",
        "from transformers import AutoTokenizer, AutoModel\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from collections import namedtuple\n",
        "import os\n",
        "# max_seq_length = 512\n",
        "\n",
        "\n",
        "# tokenizer_t_ref = AutoTokenizer.from_pretrained('allenai/specter2_base')\n",
        "# tokenize_batch_t = lambda train_text: tokenizer_t_ref.batch_encode_plus(\n",
        "#     train_text,\n",
        "#     add_special_tokens=True,\n",
        "#     padding='max_length',\n",
        "#     max_length=max_seq_length,\n",
        "#     truncation=True,\n",
        "#     return_tensors='pt',\n",
        "#     return_attention_mask=True\n",
        "# )\n",
        "\n",
        "\n",
        "\n",
        "InputExample = namedtuple('InputExample', ['texts', 'label'])\n",
        "\n",
        "\n",
        "class SIMDataset(Dataset):\n",
        "    def __init__(self, data, tokenizer,\n",
        "                 n_pairs = 0, n_pairs_random= False,\n",
        "                 emb =None, device = \"cuda\", b = 16, pttn = '* d', dir=None,\n",
        "                 map = {\"x\": \"text\", \"y\": \"label\"}):\n",
        "\n",
        "        self.b = b\n",
        "        self.device = device\n",
        "        self.tokenizer = tokenizer\n",
        "        self.emb = emb\n",
        "        self.n_pairs = n_pairs\n",
        "        self.squeeze = lambda input: {x: y.squeeze() for x, y in input.items()}\n",
        "        self.pttn = pttn\n",
        "\n",
        "        x = data[map.get(\"x\")].to_numpy()\n",
        "        y = data[map.get(\"y\")].to_numpy()\n",
        "        self.is_list = isinstance(x[0], list) or (n_pairs > 0)\n",
        "\n",
        "        # if emb:\n",
        "        #   x=x.tolist()\n",
        "        #   load = DataLoader(x, batch_size=b, shuffle=False)\n",
        "        #   with torch.no_grad(), tqdm(total=len(load), desc='Processing Batches') as pbar:\n",
        "        #     all = []\n",
        "        #     for batch_idx, d in enumerate(load):\n",
        "        #       d = self.tokenizer(d)\n",
        "        #       d = Trainer.process_batch(d, device = device)\n",
        "        #       d = self.emb(**d)\n",
        "        #       all.append(d)\n",
        "\n",
        "        #   x = torch.concat(all)\n",
        "        #   x = Trainer.process_batch(x, device = \"cpu\")\n",
        "\n",
        "\n",
        "        if n_pairs>0:\n",
        "          pairs = []\n",
        "          if not n_pairs_random:\n",
        "            for _ in range(n_pairs):\n",
        "                pairs = self.sentence_pairs_generation(x, y, pairs)\n",
        "          else:\n",
        "            pairs = self.generate_random_pairs(x, n_pairs)\n",
        "\n",
        "\n",
        "          self.examples = pairs\n",
        "        else:\n",
        "          self.examples = tuple(zip(x,y))\n",
        "\n",
        "        # print(\"len(self.examples)\")\n",
        "        # print(len(self.examples))\n",
        "\n",
        "        if emb:\n",
        "          self.examples = self.embed(emb,self.examples, dir, b)\n",
        "\n",
        "\n",
        "\n",
        "    def embed(self, emd, examples ,dir, b=18):\n",
        "      def pipe(x):\n",
        "        x = self.tokenizer(x)\n",
        "        x = Trainer.process_batch(x, device = self.device)\n",
        "        return x\n",
        "\n",
        "      # self.pttn = '* b a'\n",
        "      if dir is not None:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "      load = DataLoader(examples, batch_size=b, shuffle=False)\n",
        "\n",
        "      # emd.to(self.device)\n",
        "      with torch.no_grad(), tqdm(total=len(load), desc='Processing Batches') as pbar:\n",
        "        all = []\n",
        "        for batch_idx, d in enumerate(load):\n",
        "          x = d[0]\n",
        "          y = d[1]\n",
        "\n",
        "          if self.is_list:\n",
        "            x = [emd(**pipe(x_in)).to('cpu') for x_in in x]\n",
        "            x_k = [unpack(x_in, [[1]]*len(y), self.pttn) for x_in in x]\n",
        "            ex = tuple(zip(zip(*x_k), y))\n",
        "          else:\n",
        "            x = emd(**pipe(x))\n",
        "            ex = tuple(zip(x,y))\n",
        "\n",
        "          # Save batch data to file\n",
        "          if dir is not None:\n",
        "            batch_filename = os.path.join(dir, f'batch_{batch_idx}.pt')\n",
        "            torch.save(ex, batch_filename)\n",
        "          else:\n",
        "            all.append(ex)\n",
        "\n",
        "          x = Trainer.process_batch(x, device =\"cpu\")\n",
        "          pbar.update(1)\n",
        "\n",
        "        if dir is not None:\n",
        "          del examples\n",
        "        else:\n",
        "          if len(all) > 1:\n",
        "            all = list(chain(*all))\n",
        "          else:\n",
        "            all = all[0]\n",
        "          examples = all\n",
        "\n",
        "      torch.cuda.empty_cache()\n",
        "      return examples\n",
        "\n",
        "    def generate_random_pairs(self, texts, num_pairs=1):\n",
        "        pairs = []\n",
        "        num_texts = len(texts)\n",
        "\n",
        "        if num_pairs == 1:\n",
        "            # Generate one pair for each text with every other text\n",
        "            for i in range(num_texts):\n",
        "                for j in range(i + 1, num_texts):\n",
        "                    pairs.append(InputExample(texts=(texts[i], texts[j]), label=-1))\n",
        "        else:\n",
        "            # Generate exactly num_pairs random pairs\n",
        "            for _ in range(num_pairs):\n",
        "                pair_indices = np.random.choice(num_texts, size=2, replace=False)\n",
        "                pairs.append(InputExample(texts=(texts[pair_indices[0]], texts[pair_indices[1]]), label=-1))\n",
        "        return pairs\n",
        "\n",
        "\n",
        "\n",
        "    def sentence_pairs_generation(self, sentences, labels, pairs):\n",
        "\n",
        "        # Initialize two empty lists to hold the (sentence, sentence) pairs and\n",
        "        # labels to indicate if a pair is positive or negative\n",
        "\n",
        "        num_classes = np.unique(labels)\n",
        "        label_to_idx = {x: i for i, x in enumerate(num_classes)}\n",
        "        positive_idxs = [np.where(labels == i)[0] for i in num_classes]\n",
        "        negative_idxs = [np.where(labels != i)[0] for i in num_classes]\n",
        "\n",
        "        for first_idx in range(len(sentences)):\n",
        "            current_sentence = sentences[first_idx]\n",
        "            label = labels[first_idx]\n",
        "            second_idx = np.random.choice(positive_idxs[label_to_idx[label]])\n",
        "            positive_sentence = sentences[second_idx]\n",
        "            # Prepare a positive pair and update the sentences and labels\n",
        "            # lists, respectively\n",
        "            pairs.append(InputExample(texts=[current_sentence, positive_sentence],\n",
        "                                      label=torch.tensor([1], dtype= torch.float32)))\n",
        "\n",
        "            third_idx = np.random.choice(negative_idxs[label_to_idx[label]])\n",
        "            negative_sentence = sentences[third_idx]\n",
        "            # Prepare a negative pair of sentences and update our lists\n",
        "            pairs.append(InputExample(texts=[current_sentence, negative_sentence],\n",
        "                                      label=torch.tensor([0], dtype= torch.float32)))\n",
        "        # Return a 2-tuple of our sentence pairs and labels\n",
        "        return pairs\n",
        "\n",
        "    def get_slices(self, some_tuple, indices):\n",
        "      return tuple(some_tuple[i] for i in indices)\n",
        "\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.examples)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # batch = self.get_slices(self.examples,idx)\n",
        "        # print(type(idx))\n",
        "        # print(idx)\n",
        "        batch = self.examples[idx]\n",
        "        x,y = batch\n",
        "\n",
        "        if self.emb:\n",
        "          if self.n_pairs > 0:\n",
        "            return [x_in.squeeze() for x_in in x], y\n",
        "          else:\n",
        "            if isinstance(x, (list, tuple)):\n",
        "              return [x_in.squeeze() for x_in in x], y\n",
        "            else:\n",
        "              return x.squeeze(), y\n",
        "        else:\n",
        "          if self.n_pairs > 0:\n",
        "            if isinstance(x[0], (list, tuple)):\n",
        "              return [self.squeeze(self.tokenizer(x_in)) for x_in in x], y\n",
        "            else:\n",
        "              return [self.squeeze(self.tokenizer([x_in])) for x_in in x], y\n",
        "          else:\n",
        "            if isinstance(x, (list, tuple)):\n",
        "              return [self.squeeze(self.tokenizer([x_in])) for x_in in x], y\n",
        "            else:\n",
        "              return self.squeeze(self.tokenizer([x])), torch.Tensor([y])\n",
        "\n",
        "\n",
        "\n",
        "class LazyLoadDataset(Dataset):\n",
        "    def __init__(self, data_paths):\n",
        "        self.data_paths = data_paths\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data_paths)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # Load data at idx only when it's requested\n",
        "        tensor_path = self.data_paths[idx]\n",
        "        tensor = torch.load(tensor_path)\n",
        "        return tensor\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from collections import namedtuple\n",
        "\n",
        "# Define a namedtuple for InputExample\n",
        "InputExample = namedtuple('InputExample', ['texts', 'label'])\n",
        "\n",
        "# Define a custom dataset class\n",
        "class CustomDataset(Dataset):\n",
        "    def __init__(self, data):\n",
        "        self.data = data\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        sample = self.data[idx]\n",
        "        return sample.texts, sample.label\n",
        "\n",
        "# Define some sample data using InputExample\n",
        "sample_data = [InputExample(texts=torch.randn(3, 32, 32), label=torch.randint(0, 10, (1,)).item()) for _ in range(100)]\n",
        "\n",
        "# Create an instance of the custom dataset\n",
        "custom_dataset = CustomDataset(sample_data)\n",
        "\n",
        "# Create a DataLoader to iterate over the dataset\n",
        "batch_size = 10\n",
        "shuffle = True\n",
        "num_workers = 4\n",
        "data_loader = DataLoader(dataset=custom_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)\n",
        "\n",
        "# Iterate through the DataLoader\n",
        "for batch_idx, (texts, labels) in enumerate(data_loader):\n",
        "    # texts and labels represent a batch of data\n",
        "    print(f\"Batch {batch_idx}: Texts shape: {texts.shape}, Labels shape: {labels.shape}\")\n",
        "\n",
        "\n",
        "# sim_data = SIMDataset(df_train, tokenize_batch_t, 2)\n",
        "# train_dataloader = DataLoader(sim_data,\n",
        "#                               shuffle=True,\n",
        "#                               batch_size=batch_size)\n",
        "# for x in train_dataloader:\n",
        "#   break\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "yUrdBmR9D7e_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Similarity"
      ],
      "metadata": {
        "id": "pV26gCAQ3kaE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "from torch import nn\n",
        "from collections import namedtuple\n",
        "\n",
        "\n",
        "class SentenceEmbedding(nn.Module):\n",
        "    def __init__(self,\n",
        "                 model,\n",
        "                 normalize=False,\n",
        "                 pooling = False,\n",
        "                 head = None,\n",
        "                 **kargs):\n",
        "        super(SentenceEmbedding, self).__init__()\n",
        "\n",
        "        self.model = model\n",
        "        hidden = model.config.hidden_size\n",
        "        self.normalize = normalize\n",
        "        self.pooling = pooling\n",
        "        self.with_head = kargs.get('with_head', False)\n",
        "        self.is_bert = kargs.get('with_head', False)\n",
        "\n",
        "\n",
        "        self.head =  nn.Sequential(\n",
        "            # nn.GELU(),\n",
        "            nn.Linear(hidden, 384),\n",
        "            # nn.GELU(),\n",
        "            # nn.Linear(768, 256),\n",
        "            # nn.GELU(),\n",
        "            # nn.Linear(256, 256)\n",
        "        ) if (head is None) and self.with_head else head\n",
        "\n",
        "\n",
        "    def mean_pooling(self, model_output, attention_mask):\n",
        "        token_embeddings = model_output[0]  # First element of model_output contains all token embeddings\n",
        "        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n",
        "\n",
        "        embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n",
        "\n",
        "        return embeddings\n",
        "\n",
        "    def __getattr__(self, name: str):\n",
        "        \"\"\"Forward missing attributes to the wrapped module.\"\"\"\n",
        "        try:\n",
        "            return super().__getattr__(name)  # defer to nn.Module's logic\n",
        "        except AttributeError:\n",
        "            return getattr(self.model, name)\n",
        "\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        if isinstance(x, torch.Tensor):\n",
        "          embeddings = self.model(x)\n",
        "        else:\n",
        "          embeddings = self.model(**x)\n",
        "\n",
        "\n",
        "        if self.pooling:\n",
        "          embeddings = self.mean_pooling(embeddings,\n",
        "                                         x.get(\"attention_mask\") )\n",
        "        else:\n",
        "          embeddings = embeddings[0][:,0,:]\n",
        "\n",
        "\n",
        "        if self.with_head:\n",
        "          embeddings = self.head(embeddings)\n",
        "\n",
        "        if self.normalize:\n",
        "          embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)\n",
        "\n",
        "        return embeddings\n",
        "\n",
        "\n",
        "import torch\n",
        "from torch import nn, Tensor\n",
        "from typing import Iterable, Dict\n",
        "from typing import List\n",
        "\n",
        "\n",
        "\n",
        "class loss_sim(nn.Module):\n",
        "\n",
        "  def __init__(self, model, layer_bert = False,\n",
        "               tau = 0.1, device=\"cuda\"\n",
        "               ):\n",
        "    super(loss_sim, self).__init__()\n",
        "    self.model = model\n",
        "    self.loss = torch.nn.CosineEmbeddingLoss(margin=0.25, reduction='mean')\n",
        "    self.tau = torch.Tensor([tau]).to(device)\n",
        "    self.loss1 = torch.nn.CosineSimilarity(dim=1)\n",
        "    self.loss2 = nn.MSELoss()\n",
        "    self.n_labels = 2\n",
        "    self.loss_3 = nn.BCEWithLogitsLoss(reduction = 'mean',\n",
        "                                    pos_weight=torch.FloatTensor([2.5]))\n",
        "    # self.n_0 = torch.Tensor([0]).to('cuda')\n",
        "\n",
        "    self.layer_bert = layer_bert\n",
        "\n",
        "  def forward(self, texts, labels):\n",
        "\n",
        "    embeddings = [self.model(emd) for emd in texts]\n",
        "    labels = labels.squeeze()\n",
        "\n",
        "    s = self.loss1(embeddings[0], embeddings[1])\n",
        "\n",
        "    p = s[labels == 1]\n",
        "    n = s[labels == 0]\n",
        "\n",
        "    a = torch.mean(\n",
        "        -torch.log(\n",
        "            torch.exp(p/self.tau)/\\\n",
        "             (torch.sum(torch.exp(n/self.tau)))\n",
        "        )\n",
        "    )\n",
        "\n",
        "    # cp = torch.mean((1-torch.clamp(\n",
        "    #     p, min= 0.8\n",
        "    #     )\n",
        "    # ))\n",
        "\n",
        "\n",
        "    # cn = torch.mean((1-torch.clamp(\n",
        "    #     n, min= 0.8\n",
        "    #     )\n",
        "    # ))\n",
        "\n",
        "\n",
        "    cp = torch.mean((1-p)**2)\n",
        "    cn = torch.mean((-1-n)**2)\n",
        "    return (0.3)*a+ (0.7)*(cp+cn)\n",
        "\n",
        "    # return (0.6)*a+ (0.4)*(cp+cn)"
      ],
      "metadata": {
        "id": "w_aJVEqu3mmU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Classification"
      ],
      "metadata": {
        "id": "5RG3I0WD3nKS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # Complete model\n",
        "# class SLR_Classifier(nn.Module):\n",
        "#   def __init__(self, **data):\n",
        "#     super(SLR_Classifier, self).__init__()\n",
        "\n",
        "#     # Pre-trained model\n",
        "#     self.Encoder = data.get('model')\n",
        "\n",
        "#     self.hidden_size_input = self.Encoder.config.hidden_size\n",
        "\n",
        "#     self.hidden_size = 200\n",
        "\n",
        "\n",
        "#     # Feature Map Layer\n",
        "#     self.feature_map = nn.Sequential(\n",
        "#             # nn.LayerNorm(self.Encoder.model.config.hidden_size),\n",
        "#             # nn.BatchNorm1d(self.Encoder_config.hidden_size),\n",
        "#             # nn.Dropout(data.get(\"drop\", 0.5)),\n",
        "#             nn.Linear(self.hidden_size_input, self.hidden_size),\n",
        "#             nn.Tanh(),\n",
        "#             nn.Linear(self.hidden_size, 1),\n",
        "#             # nn.Tanh(),\n",
        "#             # nn.Linear(200, 1)\n",
        "#             # nn.Dropout(data.get(\"drop\", 0.5)),\n",
        "#         )\n",
        "\n",
        "\n",
        "#     # Initializing layer parameters\n",
        "#     # nn.init.normal_(self.feature_map[1].weight, mean=0, std=0.00001)\n",
        "#     # nn.init.zeros_(self.feature_map[1].bias)\n",
        "\n",
        "#   def __getattr__(self, name: str):\n",
        "#       \"\"\"Forward missing attributes to the wrapped module.\"\"\"\n",
        "#       try:\n",
        "#           return super().__getattr__(name)  # defer to nn.Module's logic\n",
        "#       except AttributeError:\n",
        "#           return getattr(self.model, name)\n",
        "\n",
        "#   # Feed forward\n",
        "#   def forward(self, input_ids,\n",
        "#               attention_mask,\n",
        "#               token_type_ids,\n",
        "#               train=True,\n",
        "#               **kargs):\n",
        "\n",
        "#     predict = self.Encoder(**{\"input_ids\":input_ids,\n",
        "#                               \"attention_mask\":attention_mask,\n",
        "#                               \"token_type_ids\":token_type_ids})\n",
        "#     feature = self.feature_map(predict)\n",
        "#     # logit = self.classifier(feature)\n",
        "#     representations = torch.Tensor([1])\n",
        "#     if train:\n",
        "#       return feature\n",
        "#     else:\n",
        "#       return representations, feature\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "\n",
        "# import torch\n",
        "# from torch import nn, Tensor\n",
        "# from typing import Iterable, Dict\n",
        "# from typing import List\n",
        "\n",
        "\n",
        "# class loss_slr(nn.Module):\n",
        "#   def __init__(self, model):\n",
        "#     super(loss_slr, self).__init__()\n",
        "#     self.model = model\n",
        "#     self.loss = nn.BCEWithLogitsLoss(reduction = 'mean',\n",
        "#                                     pos_weight=torch.FloatTensor([2.5])\n",
        "#                                     )\n",
        "\n",
        "\n",
        "#   def forward(self, text, label, **kargs):\n",
        "#     embeddings = self.model(**text)\n",
        "#     return self.loss(embeddings.to(torch.float), label.to(torch.float))\n"
      ],
      "metadata": {
        "id": "W2IqXyrX3pel"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Distill"
      ],
      "metadata": {
        "id": "G-nYzzaINwbq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Criar codigo para que consiga independente da quantidade de textos, fazer o embeding e retornar da mesma forma.\n",
        "\n",
        "```\n",
        "[\n",
        "  ((text1,...,), label),\n",
        "  ...,\n",
        "  ((text1,...,), label),\n",
        "]\n",
        "```\n",
        "\n",
        "```\n",
        "[\n",
        "  ((emb1,...,), label),\n",
        "  ...,\n",
        "  ((emb1,...,), label),\n",
        "]\n",
        "```\n",
        "\n",
        "deve ser squezaded o emb, pode ser passado o batch para agregar quantidade de exemplos q pode ser passado"
      ],
      "metadata": {
        "id": "vZZs0PJsYwso"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import DataLoader\n",
        "from einops import pack, unpack\n",
        "\n",
        "n = 50\n",
        "ex = 2\n",
        "input_data = torch.randn(n,ex, 5)\n",
        "output_data = torch.randn(n, 1)\n",
        "batch = 20\n",
        "\n",
        "\n",
        "def batch_load(input, output, batch):\n",
        "  input_data = input.get('data')\n",
        "  output_data = output.get('data')\n",
        "\n",
        "  input_cfg = input.get('config')\n",
        "  output_cfg = output.get('config')\n",
        "\n",
        "\n",
        "  full = input_data.shape[0]//batch\n",
        "  batch_as = [[batch]]*full\n",
        "  rest =  input_data.shape[0]%batch\n",
        "\n",
        "  if rest > 0:\n",
        "    batch_as += [[rest]]\n",
        "\n",
        "  # input_data.shape\n",
        "  in_a = unpack(input_data, batch_as, input_cfg)\n",
        "  out_a = unpack(output_data, batch_as, output_cfg)\n",
        "\n",
        "\n",
        "  return zip(in_a, out_a)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "a = batch_load(\n",
        "    input = {\n",
        "        \"data\": input_data,\n",
        "        \"config\": \"* ex x\"\n",
        "    },\n",
        "    output = {\n",
        "        \"data\": output_data,\n",
        "        \"config\": \"* y\"\n",
        "    },\n",
        "    batch = batch\n",
        ")\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "4pIjfkhpyuUa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "##### similatiry distill\n",
        "# the model is trained to distill the similarity of the teacher\n",
        "\n",
        "class few_shot_loss_distill(nn.Module):\n",
        "  def __init__(self, model_s, model_t,\n",
        "               t=0, loss=0):\n",
        "    super(few_shot_loss_distill, self).__init__()\n",
        "    self.model_t = model_t\n",
        "    self.model_t.eval()\n",
        "    self.model_s = model_s\n",
        "    self.loss_choice = loss\n",
        "    self.tau = 0.065\n",
        "    self.t = t\n",
        "    self.loss1 = torch.nn.CosineSimilarity(dim=1)\n",
        "    self.distill = lambda x: F.log_softmax(x, dim = 0)\n",
        "\n",
        "    self.loss = torch.nn.CosineEmbeddingLoss(margin=0.25, reduction='mean')\n",
        "    # self.proj = nn.Sequential(\n",
        "    #         # nn.GELU(),\n",
        "    #         nn.Linear(384, 384),\n",
        "    #         # nn.GELU(),\n",
        "    #         # nn.Linear(384, 384)\n",
        "    #         )\n",
        "\n",
        "\n",
        "    # # return self.loss(embeddings[0],embeddings[1],labels)\n",
        "\n",
        "\n",
        "  def forward(self, texts, labels):\n",
        "    # print(\"texts\")\n",
        "    # global texts\n",
        "    texts = tuple(texts)\n",
        "    # print(texts)\n",
        "    # embeddings_s = [self.model_s(**emd) for emd in texts]\n",
        "    # embeddings_s2 = [self.proj(self.model_s(**emd)) for emd in texts]\n",
        "    embeddings_s = [self.model_s(**emd) for emd in texts]\n",
        "\n",
        "    # if self.model_t:\n",
        "    embeddings_t = [self.model_t(**emd) for emd in texts]\n",
        "\n",
        "\n",
        "    # embeddings_s = torch.concat(embeddings_s)\n",
        "    # embeddings_t = torch.concat(embeddings_t)\n",
        "\n",
        "    # sz = self.loss1(embeddings_s,embeddings_t)\n",
        "\n",
        "    # embeddings = pack([embeddings_t, embeddings_s], '* b h')[0]\n",
        "    # loss_kd =-self.distill(embeddings/self.tau)\n",
        "    # return torch.mean(loss_kd)\n",
        "\n",
        "    # print(sz.shape)\n",
        "    # labels = labels.squeeze()\n",
        "    # labels[labels == 0] = -1\n",
        "\n",
        "    # p = sz[labels == 1]\n",
        "    # n = sz[labels == -1]\n",
        "\n",
        "\n",
        "    # return sz\n",
        "    # return self.loss(embeddings[0],embeddings[1],labels)\n",
        "\n",
        "    s_t = self.loss1(*embeddings_t)\n",
        "    s_s = self.loss1(*embeddings_s)\n",
        "\n",
        "    labels[labels == 0] = -1\n",
        "    labels = labels.squeeze()\n",
        "\n",
        "    p_t = s_t[labels == 1]\n",
        "    n_t = s_t[labels == -1]\n",
        "\n",
        "    p_s = s_s[labels == 1]\n",
        "    n_s = s_s[labels == -1]\n",
        "\n",
        "    s = (s_t-s_s)\n",
        "\n",
        "    p = s[labels == 1]\n",
        "    n = s[labels == -1]\n",
        "\n",
        "    if self.loss_choice == 0:\n",
        "      # s = s[s>]\n",
        "      dist1 = torch.mean(\n",
        "          -torch.log(\n",
        "              torch.exp(s/self.tau)/\\\n",
        "              (torch.sum(torch.exp(s/self.tau)))\n",
        "          )\n",
        "      )\n",
        "\n",
        "      return  dist1\n",
        "\n",
        "    elif self.loss_choice == 1:\n",
        "      dist = torch.mean(\n",
        "          -torch.log(\n",
        "              torch.exp(p/self.tau)/\\\n",
        "              (torch.sum(torch.exp(n/self.tau)))\n",
        "          )\n",
        "      )\n",
        "      return  dist\n",
        "\n",
        "\n",
        "    # dist_s = torch.mean(\n",
        "    #     -torch.log(\n",
        "    #         torch.exp(p_s/self.tau)/\\\n",
        "    #          (torch.sum(torch.exp(n_s/self.tau)))\n",
        "    #     )\n",
        "    # )\n",
        "\n",
        "    # dist_t = torch.mean(\n",
        "    #     -torch.log(\n",
        "    #         torch.exp(p_t/self.tau)/\\\n",
        "    #          (torch.sum(torch.exp(n_t/self.tau)))\n",
        "    #     )\n",
        "    # )\n",
        "\n",
        "\n",
        "\n",
        "    # return torch.mean((s_t-s_s)**2)\n",
        "\n",
        "\n",
        "\n",
        "##### head distill\n",
        "# make a new layer to distill the logits of the teacher\n",
        "\n",
        "class loss_distill_to_light(nn.Module):\n",
        "  def __init__(self, model_s, model_t = None,\n",
        "               tau = 0.1, t=0):\n",
        "    super(loss_distill_to_light, self).__init__()\n",
        "    self.model_t = model_t\n",
        "    self.model_s = model_s\n",
        "\n",
        "    self.tau = tau\n",
        "    self.t = t\n",
        "\n",
        "    # (workers batch hidden)\n",
        "    self.loss1 = torch.nn.CosineSimilarity(dim=1)\n",
        "    self.distill = lambda x: F.log_softmax(x, dim = self.t)\n",
        "\n",
        "\n",
        "  def forward(self, texts, labels):\n",
        "    if self.model_t:\n",
        "      embeddings_t = [self.model_t(**emd) for emd in texts]\n",
        "    embeddings_s = [self.model_s(**emd) for emd in texts]\n",
        "\n",
        "    embeddings = pack([embeddings_t, embeddings_s], '* b h')[0]\n",
        "\n",
        "    loss_kd =-self.distill(embeddings/self.tau)\n",
        "\n",
        "    return torch.mean(loss_kd, self.t)\n",
        "\n",
        "    # s = self.loss1(embeddings_t, embeddings_s)\n",
        "\n",
        "    # return torch.sum(s**2)\n",
        "\n",
        "\n",
        "\n",
        "# import itertools\n",
        "# a = itertools.tee(a,2)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "i-dr-4TINzdz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 📓 Train pipeline\n"
      ],
      "metadata": {
        "id": "vVUXhbqf3sVw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Class"
      ],
      "metadata": {
        "id": "6WIel2kEUaeM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import gc\n",
        "\n",
        "\n",
        "\n",
        "class Trainer:\n",
        "    def __init__(self,\n",
        "                 losses=None,\n",
        "                 optimizer=None,\n",
        "                 train_dataloader=None,\n",
        "                 total_iterations=None,\n",
        "                 val_dataloader = None,\n",
        "                 lr=None,\n",
        "                 accumulation_steps=None,\n",
        "                 scheduler= None,\n",
        "                 device = \"cpu\",\n",
        "                 **kwargs):\n",
        "\n",
        "        self.losses = losses\n",
        "        self.train_dataloader = train_dataloader\n",
        "        self.val_dataloader = val_dataloader\n",
        "        self.total_iterations = total_iterations\n",
        "        self.lr = lr\n",
        "        self.accumulation_steps = accumulation_steps\n",
        "        self.optimizer = optimizer\n",
        "        self.scheduler = scheduler\n",
        "\n",
        "\n",
        "        self.device =device #= \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "        self.loss_values = []\n",
        "\n",
        "\n",
        "\n",
        "    @classmethod\n",
        "    def process_batch(cls, batch, device = 'cuda', generator = True):\n",
        "\n",
        "        def to_device_recursive(item):\n",
        "            if isinstance(item, dict):\n",
        "                return {key: to_device_recursive(value) for key, value in item.items()}\n",
        "            elif isinstance(item, (list, tuple)):\n",
        "                if generator:\n",
        "                  return (to_device_recursive(value) for value in item)\n",
        "                else:\n",
        "                  return [to_device_recursive(value) for value in item]\n",
        "            else:\n",
        "                return item.to(device)\n",
        "\n",
        "        batch = to_device_recursive(batch)\n",
        "        return batch\n",
        "\n",
        "\n",
        "    def compute_loss(self, batch):\n",
        "        batch = self.process_batch(batch, self.device)\n",
        "        # y = self.process_batch(batch[1], self.device)\n",
        "\n",
        "        return self.losses(*batch)\n",
        "\n",
        "\n",
        "    def gradient_update(self):\n",
        "        if (self.iteration + 1) % self.accumulation_steps == 0:\n",
        "            self.optimizer.step()\n",
        "            if self.scheduler:\n",
        "              self.scheduler.step()\n",
        "\n",
        "            self.optimizer.zero_grad()\n",
        "\n",
        "\n",
        "    def train(self, losses=None,\n",
        "                 optimizer=None,\n",
        "                 train_dataloader=None,\n",
        "                 total_iterations=None,\n",
        "                 val_dataloader = None,\n",
        "                 lr=None,\n",
        "                 accumulation_steps=None,\n",
        "                 scheduler= None,\n",
        "                 device = \"cpu\",\n",
        "                 **kwargs):\n",
        "\n",
        "        train_dataloader = self.train_dataloader if train_dataloader is None else train_dataloader\n",
        "        val_dataloader = self.val_dataloader if val_dataloader is None else val_dataloader\n",
        "        total_iterations = self.total_iterations if total_iterations is None else total_iterations\n",
        "        accumulation_steps = self.accumulation_steps if accumulation_steps is None else accumulation_steps\n",
        "        total_iterations = self.total_iterations if total_iterations is None else total_iterations\n",
        "\n",
        "        self.losses = self.losses if losses is None else losses\n",
        "        self.lr = self.lr if lr is None else lr\n",
        "        self.optimizer = self.optimizer if optimizer is None else optimizer\n",
        "        self.scheduler = self.scheduler if scheduler is None else scheduler\n",
        "\n",
        "        pbar = tqdm(total=total_iterations, desc='Training', unit='batch', position=0)\n",
        "\n",
        "        batch_maker = iter(train_dataloader)\n",
        "        for self.iteration in range(total_iterations):\n",
        "\n",
        "            # indo para o proximo batch\n",
        "            try:\n",
        "              batch = next(batch_maker)\n",
        "            except StopIteration:\n",
        "              batch = next(iter(train_dataloader))\n",
        "\n",
        "            # computa a funcao de perda e os gradientes\n",
        "            loss = self.compute_loss(batch)\n",
        "\n",
        "            del batch\n",
        "\n",
        "            loss = loss / self.accumulation_steps\n",
        "            loss.backward()\n",
        "\n",
        "            # atualiza os parametros\n",
        "            self.gradient_update()\n",
        "\n",
        "            # if self.val_dataloader:\n",
        "              # self.validate()\n",
        "\n",
        "            pbar.set_postfix({\n",
        "                'Loss': loss.item(),\n",
        "                # 'lr': self.lr.item()\n",
        "                })\n",
        "            self.loss_values.append(loss.item())\n",
        "\n",
        "            pbar.update(1)\n",
        "\n",
        "\n",
        "        gc.collect()\n",
        "        torch.cuda.empty_cache()\n",
        "\n",
        "        # Adicione sua lógica aqui para calcular métric as ou fazer avaliações durante o treinamento\n",
        "        # Certifique-se de fechar a barra de progresso após o treinamento\n",
        "        pbar.close()\n",
        "        self.plot_loss_curve()\n",
        "\n",
        "\n",
        "\n",
        "    def plot_loss_curve(self):\n",
        "        plt.figure(figsize=(10, 6))\n",
        "        plt.plot(self.loss_values, label='Training Loss')\n",
        "        plt.xlabel('Iteration')\n",
        "        plt.ylabel('Loss')\n",
        "        plt.title('Training Loss Curve')\n",
        "        plt.legend()\n",
        "        plt.show()\n",
        "\n",
        "\n",
        "\n",
        "# #############################\n",
        "\n",
        "\n",
        "\n",
        "import torch\n",
        "import gc\n",
        "\n",
        "import torch\n",
        "import gc\n",
        "from tqdm import tqdm\n",
        "\n",
        "def get_batch_labels(model, dataloader, device=\"cuda\"):\n",
        "    preds = []\n",
        "    labels = []\n",
        "\n",
        "    model.eval()\n",
        "    with torch.no_grad(), tqdm(total=len(dataloader), desc='Processing Batches') as pbar:\n",
        "        for batch in dataloader:\n",
        "            # print(batch)\n",
        "            x = Trainer.process_batch(batch[0], device)\n",
        "            pred = model(x)\n",
        "            x = Trainer.process_batch(x, 'cpu')\n",
        "            pred = Trainer.process_batch(pred, 'cpu')\n",
        "\n",
        "            preds.append(pred)\n",
        "            labels.append(batch[1])\n",
        "\n",
        "            del x\n",
        "\n",
        "            pbar.update(1)\n",
        "\n",
        "        gc.collect()\n",
        "        torch.cuda.empty_cache()\n",
        "\n",
        "    return preds, labels\n",
        "\n",
        "# Exemplo de chamada da função\n",
        "# Substitua 'sua_funcao_que_retorna_as_predicoes' pelo nome da função em seu modelo\n",
        "# que retorna as predições desejadas.\n",
        "# preds, labels = get_batch_labels(model_s2, train_dataloader, Trainer)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def configure_optimizer(self, model, lr):\n",
        "    # Função para configurar o otimizador\n",
        "    self.optimizer = Adam(model.parameters(), lr=lr)\n",
        "    self.scheduler = LambdaLR(self.optimizer, lr_lambda=lambda epoch: 0.7 ** (epoch // 20) if 0.7 ** (epoch // 20) > 1e-2 else 0.05)\n",
        "\n",
        "\n",
        "def simulate_and_plot_lr_schedule(self, total_iterations_sim=None, lr = None):\n",
        "    dumb_model = nn.Linear(1, 1)\n",
        "    if lr:\n",
        "      # Configurando o otimizador usando a função configuradora\n",
        "      self.configure_optimizer(dumb_model, lr=lr)\n",
        "    else:\n",
        "      self.configure_optimizer(dumb_model, lr=self.lr)\n",
        "\n",
        "    if total_iterations_sim:\n",
        "      total_iterations_sim = total_iterations_sim\n",
        "    else:\n",
        "      total_iterations_sim = self.total_iterations\n",
        "\n",
        "    lr_values = []\n",
        "    for self.iteration in range(total_iterations_sim):\n",
        "      self.gradient_update()\n",
        "      current_lr = self.optimizer.param_groups[0]['lr']\n",
        "      lr_values.append(current_lr)\n",
        "\n",
        "    # default padron\n",
        "    self.configure_optimizer(self.model, self.lr)\n",
        "\n",
        "\n",
        "    # Plotando os valores do learning rate\n",
        "    plt.plot(lr_values)\n",
        "    plt.xlabel('Iteration')\n",
        "    plt.ylabel('Learning Rate')\n",
        "    plt.title('Learning Rate Schedule during Training Simulation')\n",
        "    plt.show()\n",
        "    return lr_values\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "eyze5LnI9HCE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## metrics"
      ],
      "metadata": {
        "id": "vf9HASWycrIO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "from sklearn.datasets import make_classification\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import roc_curve\n",
        "import matplotlib.pyplot as plt\n",
        "from sklearn.metrics import confusion_matrix\n",
        "from sklearn.metrics import f1_score\n",
        "from scipy.stats import gaussian_kde\n",
        "\n",
        "import matplotlib as mpl\n",
        "# mpl.rcParams['hatch.linewidth'] = 0.1  # previous pdf hatch linewidth\n",
        "mpl.rcParams['hatch.linewidth'] = 5.0  # previous svg hatch linewidth\n",
        "\n",
        "def plot_density(data, color, label, reverse = False):\n",
        "    # Sort the data\n",
        "    data_sorted = np.sort(data)\n",
        "    print((data_sorted.min(), data_sorted.max()),label)\n",
        "\n",
        "    ##### estimative of cumulative density\n",
        "    # Estimate the density values\n",
        "    # print(\"data\")\n",
        "    # print(data_sorted)\n",
        "    kde = gaussian_kde(data_sorted)\n",
        "    ym = kde(data_sorted)\n",
        "\n",
        "    # Normalize y-values between 0 and 1\n",
        "    ym_normalized = (ym - ym.min()) / (ym.max() - ym.min())\n",
        "\n",
        "    zero = np.array([0])\n",
        "\n",
        "    # Extend data_sorted and ym_normalized with zeros at the beginning and end\n",
        "    data_sorted_extended = np.concatenate([data_sorted[:1], data_sorted, data_sorted[-1:]])\n",
        "    ym_normalized_extended = np.concatenate([zero, ym_normalized, zero])\n",
        "    ym = np.concatenate([zero, ym, zero])\n",
        "    #####    #####    #####\n",
        "\n",
        "    # Plot the density estimates with normalized y-values\n",
        "    plt.plot(data_sorted_extended, ym_normalized_extended, label=f\"Densidade {label}\", color=color)\n",
        "    plt.fill_between(data_sorted_extended, ym_normalized_extended, color=color, alpha=0.3)\n",
        "\n",
        "    if reverse:\n",
        "      cum = 1-np.cumsum(ym)/np.sum(ym)\n",
        "    else:\n",
        "      cum = np.cumsum(ym)/np.sum(ym)\n",
        "\n",
        "    plt.plot(data_sorted_extended, cum, label=f\"Acumulada {label}\", color=color, linestyle='--')\n",
        "    plt.fill_between(data_sorted_extended, cum, color=color, alpha=0.3,  hatch='*')\n",
        "    plt.ylabel('Densidade', fontsize=17, fontweight='bold')\n",
        "    plt.xlabel('Valor', fontsize=17, fontweight='bold')\n",
        "    plt.yticks(fontsize=14)\n",
        "    plt.xticks(fontsize=14)\n",
        "\n",
        "    plt.grid(True)\n",
        "\n",
        "\n",
        "# Função para calcular WSS com base em um recall desejado\n",
        "def metrics(probs, labels, recall_target=0.95):\n",
        "    fpr, tpr, thresholds = roc_curve(labels, probs)\n",
        "\n",
        "    # Encontrar o índice correspondente ao recall desejado\n",
        "    idx_target_recall = sum(tpr < recall_target)\n",
        "\n",
        "    # Obter o limiar correspondente ao recall desejado\n",
        "    threshold_target_recall = thresholds[idx_target_recall]\n",
        "    predict_trash = (probs > threshold_target_recall).astype(int)\n",
        "\n",
        "\n",
        "    # probs[idx_target_recall:]\n",
        "    # Prever com base no limiar\n",
        "    # predict_trash = torch.sigmoid(logit).squeeze() >= threshold_target_recall\n",
        "\n",
        "    # Calcular a matriz de confusão\n",
        "    CM = confusion_matrix(labels, predict_trash)\n",
        "    tn, fp, fne, tp = CM.ravel()\n",
        "\n",
        "    # Calcular as métricas\n",
        "    P = tp + fne\n",
        "    N = tn + fp\n",
        "    recall = tp / P\n",
        "    wss = (tn + fne) / len(labels) - (1 - recall)\n",
        "    awss = (tn / N - fne / P)\n",
        "    f1 = f1_score(predict_trash, predict_trash)\n",
        "\n",
        "\n",
        "    # print(predict_trash)\n",
        "\n",
        "    return {\n",
        "        \"wss\": round(wss, 4),\n",
        "        \"awss\": round(awss, 4),\n",
        "        \"R\": round(recall, 4),\n",
        "        \"f1\": round(f1, 4),\n",
        "        \"CM\": CM,\n",
        "        'prox.': np.mean((predict_trash == labels)).astype(float)\n",
        "    }\n",
        "\n",
        "import torch\n",
        "from sklearn.metrics import roc_curve, auc, confusion_matrix\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "def plot_metrics(y_scores, y_true, recall_target=0.95, save=None):\n",
        "    # Calcula a curva ROC\n",
        "    fpr, tpr, thresholds = roc_curve(y_true, y_scores)\n",
        "    roc_auc = auc(fpr, tpr)\n",
        "\n",
        "    # Plota a curva ROC\n",
        "    plt.figure(figsize=(12, 5))\n",
        "    plt.subplot(1, 2, 1)\n",
        "    plt.plot(fpr, tpr, color='darkorange', lw=2, label='cuvra ROC (area = {:.2f})'.format(roc_auc))\n",
        "    plt.axhline(y=recall_target, color='red', linestyle='--', label=f'Recall esperado ({recall_target})')\n",
        "\n",
        "    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')\n",
        "    plt.xlabel('Falso Positivo (%)', fontweight='bold', fontsize=14)\n",
        "    plt.ylabel('Verdadeiro Positivo (%)', fontweight='bold', fontsize=14)\n",
        "    # plt.title('Receiver Operating Characteristic (ROC) Curve')\n",
        "    legend = plt.legend(loc='lower right', fontsize=10)\n",
        "\n",
        "    for text in legend.get_texts():\n",
        "      text.set_fontweight('bold')\n",
        "\n",
        "    # Calcula a matriz de confusão e outras métricas\n",
        "    y_hat = (y_scores > 0.5)\n",
        "    # y_true = y_true\n",
        "\n",
        "    CM = confusion_matrix(y_true, y_hat)\n",
        "    tn, fp, fne, tp = CM.ravel()\n",
        "\n",
        "    # Calcular as métricas\n",
        "    P = tp + fne\n",
        "    N = tn + fp\n",
        "    recall = tp / P\n",
        "    wss = (tn + fne) / len(y_true) - (1 - recall)\n",
        "    awss = (tn / N - fne / P)\n",
        "\n",
        "    # Assuming y_true and y_pred are your ground truth and predicted labels\n",
        "    f1 = f1_score(y_true, y_hat)\n",
        "\n",
        "\n",
        "    # Plot a matriz como cores usando um mapa de calor\n",
        "    plt.subplot(1, 2, 2)\n",
        "    m_normalized = np.divide(CM, CM.sum(axis=1, keepdims=True))\n",
        "    plt.imshow(m_normalized, cmap='viridis', interpolation='nearest')\n",
        "    plt.colorbar()\n",
        "    # plt.title('Confusion Matrix as Colors')\n",
        "\n",
        "    plt.tight_layout()\n",
        "\n",
        "    if save is not None:\n",
        "        plt.savefig(f'{save}/plt1.pdf')  # Salva como PNG\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "    plt.figure(figsize=(12, 5))\n",
        "    plot_density(y_scores[y_true == 0], 'blue', '-1', reverse=True)\n",
        "    plot_density(y_scores[y_true == 1], 'orange', '1')\n",
        "    plt.axvline(x=0.5, color='red', linestyle='--')\n",
        "    # Set the axis limits\n",
        "    plt.xlim(0, 1)\n",
        "    plt.ylim(0, 1)\n",
        "    # Add labels and title\n",
        "    # plt.title('Normalized Probability Density Estimation')\n",
        "\n",
        "    # Show legend\n",
        "    legend = plt.legend(fontsize=12, loc='upper center')\n",
        "\n",
        "    for text in legend.get_texts():\n",
        "      text.set_fontweight('bold')\n",
        "\n",
        "    if save is not None:\n",
        "        plt.savefig(f'{save}/plt2.pdf')  # Salva como PNG\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "    return {\n",
        "     \"95\": metrics(y_scores, y_true),\n",
        "     'observed': {\n",
        "        \"wss\": round(wss, 4),\n",
        "        \"awss\": round(awss, 4),\n",
        "        \"f1\": round(f1, 4),\n",
        "        \"R\": round(recall, 4),\n",
        "        \"CM\": CM,\n",
        "        \"roc_auc\":roc_auc,\n",
        "        'prox.': np.mean(y_hat == y_true).astype(float),\n",
        "        \"fpr\":fpr,\n",
        "        \"tpr\":tpr,\n",
        "    }\n",
        "    }\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def plot_rocs(x_list, names):\n",
        "    plt.figure(figsize=(12, 5))\n",
        "    for x, name in zip(x_list, names):\n",
        "      fpr = x.get(\"fpr\")\n",
        "      tpr = x.get(\"tpr\")\n",
        "      roc_auc = x.get(\"roc_auc\")\n",
        "      plt.plot(fpr, tpr, lw=2, label=f'{name} (área = {roc_auc:.2f})')\n",
        "\n",
        "    plt.axhline(y=0.95, color='red', linestyle='--', label=f'Recall 95%')\n",
        "    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')\n",
        "    plt.xlabel('Falso Positivo (%)')\n",
        "    plt.ylabel('Verdadeiro Positivo (%)')\n",
        "    # plt.title('Curva ROC')\n",
        "    plt.legend(loc='lower right')\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def plot_cosine_similarity_matrix(similarities, vmin=-1.0, vmax=1.0):\n",
        "    \"\"\"\n",
        "    Plot the cosine similarity matrix with a horizontal colorbar.\n",
        "\n",
        "    Args:\n",
        "    - similarities (torch.Tensor): Cosine similarity matrix\n",
        "    - vmin (float): Minimum value for colormap (default: -1.0)\n",
        "    - vmax (float): Maximum value for colormap (default: 1.0)\n",
        "    \"\"\"\n",
        "    # plt.figure(figsize=(8, 6))\n",
        "    plt.figure(figsize=(10, 10))\n",
        "\n",
        "    plt.imshow(similarities.numpy(), interpolation='none',\n",
        "                vmin=vmin, vmax=vmax)\n",
        "\n",
        "    # colorbar = plt.colorbar(label='Similaridade de Cosseno')\n",
        "    # colorbar.ax.tick_params(labelsize=20)  # Set font size for colorbar ticks\n",
        "    # colorbar.ax.set_ylabel('Similaridade de Cosseno', fontsize=22, weight=\"bold\")  # Set font size and weight for colorbar label\n",
        "\n",
        "       # plt.title('Matriz de similaridade de cossenos')\n",
        "    plt.xlabel('Índice', fontsize=20, weight=\"bold\")\n",
        "    plt.ylabel('Índice', fontsize=20, weight=\"bold\")\n",
        "    plt.yticks(fontsize=18)\n",
        "    plt.xticks(fontsize=18)\n",
        "\n",
        "\n",
        "    plt.show()\n",
        "\n"
      ],
      "metadata": {
        "id": "QRGzx6l4K4Ho"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import numpy as np\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.manifold import TSNE\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def evaluate_model(model, train_dataloader, test_dataloader, device, dtype=torch.float32):\n",
        "    emb_test_preds = get_batch_labels(model, test_dataloader, device)  # Assuming Trainer is defined\n",
        "    emb_train_preds = get_batch_labels(model, train_dataloader, device)\n",
        "\n",
        "    emb_test_preds = [\n",
        "        torch.cat(emb_test_preds[0]).to(\"cpu\").squeeze().to(dtype=dtype),\n",
        "        torch.cat(emb_test_preds[1]).to(\"cpu\").squeeze().to(dtype=dtype)\n",
        "    ]\n",
        "\n",
        "    emb_train_preds = [\n",
        "        torch.cat(emb_train_preds[0]).to(\"cpu\").squeeze().to(dtype=dtype),\n",
        "        torch.cat(emb_train_preds[1]).to(\"cpu\").squeeze().to(dtype=dtype)\n",
        "    ]\n",
        "\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "    return emb_test_preds, emb_train_preds\n",
        "\n",
        "\n",
        "def visualize_tsne(representations, labels,\n",
        "                   n_components=2,\n",
        "                   perplexity=30, save=None, alpha=1):\n",
        "    tsne = TSNE(n_components=n_components, perplexity=perplexity, random_state=42)\n",
        "    embeddings_2d = tsne.fit_transform(representations)\n",
        "\n",
        "    plt.figure(figsize=(10, 10))\n",
        "\n",
        "    for i in labels.unique():\n",
        "        indices = labels == i\n",
        "        plt.scatter(embeddings_2d[indices, 0],\n",
        "                    embeddings_2d[indices, 1], label=f'Classe {int(i)}', alpha=alpha)\n",
        "\n",
        "    plt.yticks(fontsize=12)\n",
        "    plt.xticks(fontsize=12)\n",
        "\n",
        "    # plt.title(f't-SNE do Model Embeddings ({n_components} components)')\n",
        "    legend = plt.legend(fontsize=14)\n",
        "\n",
        "    for text in legend.get_texts():\n",
        "      text.set_fontweight('bold')\n",
        "\n",
        "\n",
        "    if save is not None:\n",
        "      plt.savefig(f'{save}/tsne.pdf')\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "def logistic_regression_probabilities(Train, Test= None, random_state=None):\n",
        "    \"\"\"\n",
        "    Realiza a regressão logística e retorna as probabilidades de pertencer à classe 1.\n",
        "\n",
        "    Parâmetros:\n",
        "    - features: Um array-like bidimensional (n, k) contendo as características dos exemplos.\n",
        "    - categories: Uma lista de categorias (0, 1) para cada exemplo.\n",
        "    - test_size: A proporção do conjunto de dados a ser utilizado como conjunto de teste.\n",
        "    - random_state: Semente para controle de reprodutibilidade.\n",
        "\n",
        "    Retorna:\n",
        "    - Uma lista de probabilidades para pertencer à classe 1 para cada exemplo.\n",
        "    \"\"\"\n",
        "    # Dividindo o conjunto de dados em treino e teste\n",
        "    # X_train, X_test, y_train, y_test = train_test_split(features, categories, test_size=test_size, random_state=random_state)\n",
        "    preds, label = Train\n",
        "    # Criando e treinando o modelo de regressão logística\n",
        "    model = LogisticRegression(max_iter=1000)\n",
        "    model.fit(preds, label)\n",
        "\n",
        "    # Obtendo as probabilidades no conjunto de teste\n",
        "    probabilities = model.predict_proba(preds)[:, 1]\n",
        "\n",
        "    if Test != None:\n",
        "      return model.predict_proba(Test)[:, 1]\n",
        "\n",
        "    return probabilities\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def eval_pipeline(model, train_dataloader, test_dataloader, device=\"cuda\", perplexity=30, save=None, tsne=False):\n",
        "\n",
        "\n",
        "    emb_test_preds, emb_train_preds = evaluate_model(model, train_dataloader,\n",
        "                                                     test_dataloader, device)\n",
        "    if tsne:\n",
        "      visualize_tsne(emb_test_preds[0], emb_test_preds[1], perplexity =perplexity, save=save)\n",
        "\n",
        "\n",
        "\n",
        "    probs = logistic_regression_probabilities(\n",
        "        Train = emb_train_preds,\n",
        "        Test = emb_test_preds[0])\n",
        "\n",
        "    # print(emb_test_preds[1].numpy())\n",
        "    res = plot_metrics(probs,\n",
        "                        emb_test_preds[1].numpy(),\n",
        "                        recall_target=0.95,\n",
        "                        save=save)\n",
        "    return res\n",
        "\n"
      ],
      "metadata": {
        "id": "GrCFgYOZiZOn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Utils"
      ],
      "metadata": {
        "id": "JGLOTDg2FxlB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.optim import Adam\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from tqdm import tqdm\n",
        "from torch.optim.lr_scheduler import LambdaLR\n",
        "######\n",
        "\n",
        "##########################################################################\n",
        "\n",
        "\n",
        "def main_pipe_train(smodel, submod, df_train,train_dataloader, train_dataloader_eval, test_dataloader_eval, config, save=None):\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "    # Mod(smodel_mod_t)\n",
        "    Mod(smodel)\n",
        "    smodel.util.map(\n",
        "        name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "    # activate in the last layers\n",
        "    Mod(submod)\n",
        "    submod.util.activate(\n",
        "        layer_names= [\"query\", 'value'],\n",
        "        key = 'set',\n",
        "        config = {\n",
        "            'name':'test',\n",
        "            **config\n",
        "            })\n",
        "\n",
        "\n",
        "    smodel.util.trainable_params()\n",
        "    smodel = smodel.to(smodel.dtype)\n",
        "\n",
        "    def lr_lambda(epoch):\n",
        "        return 1 - (epoch / total_iterations)**2\n",
        "\n",
        "    n_examples = len(df_train)*config[\"n_pairs\"]\n",
        "\n",
        "    interation = int(n_examples/config[\"batch_size\"])\n",
        "    total_iterations = config[\"n_epocs\"]*(interation)\n",
        "\n",
        "\n",
        "    # Create the linear learning rate scheduler\n",
        "    optimizer = Adam(smodel.parameters(), lr=config[\"lr\"])\n",
        "    scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)\n",
        "\n",
        "    #### train\n",
        "    # smodel\n",
        "    trainer = Trainer(losses=loss_sim(smodel, tau = config[\"tau\"], device=\"cuda\").to(\"cuda\"),\n",
        "                      optimizer=optimizer,\n",
        "                      scheduler=scheduler,\n",
        "                      accumulation_steps = config[\"accumulation_steps\"],\n",
        "                      train_dataloader=train_dataloader,\n",
        "                      total_iterations=total_iterations,\n",
        "                      device = 'cuda')\n",
        "\n",
        "\n",
        "    print(f\"Quantidade de amostras\\n {n_examples}\")\n",
        "    print(f\"Quantidade de exemplos\\n {len(df_train)}\")\n",
        "\n",
        "\n",
        "\n",
        "    w = smodel.model.encoder.layer[-1].attention.self.value.Mix['Adapt'].loraA\n",
        "\n",
        "    print(submod)\n",
        "    print()\n",
        "    print(w)\n",
        "\n",
        "    trainer.train()\n",
        "\n",
        "    print(w)\n",
        "    smodel.to('cpu')\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "    from pprint import pprint\n",
        "    ##########################################################################\n",
        "\n",
        "    torch.cuda.empty_cache()\n",
        "    smodel.cuda()\n",
        "    res = eval_pipeline(smodel,\n",
        "                        train_dataloader=train_dataloader_eval,\n",
        "                        test_dataloader=test_dataloader_eval,\n",
        "                  device=\"cuda\", save=save)\n",
        "    smodel.cpu()\n",
        "\n",
        "\n",
        "    pprint(res)\n",
        "\n",
        "\n",
        "    torch.cuda.empty_cache()\n",
        "    return res, trainer\n",
        "\n",
        "\n",
        "def eval_model_emb(smodel, lod):\n",
        "  smodel.cuda()\n",
        "  preds= []\n",
        "  labels= []\n",
        "  for i,lo in enumerate(lod):\n",
        "    emb_test, emb_train = evaluate_model(smodel,*lo, \"cuda\")\n",
        "    # print(emb_test[1])\n",
        "    i = i+1\n",
        "    emb_test[1][emb_test[1]==1] = i\n",
        "    emb_test[1][emb_test[1]==0] = -i\n",
        "    # print(emb_test[1])\n",
        "\n",
        "    preds.append(emb_test[0])\n",
        "    labels.append(emb_test[1])\n",
        "\n",
        "  smodel.cpu()\n",
        "  torch.cuda.empty_cache()\n",
        "  if len(lod)>1:\n",
        "    preds = torch.concat(preds)\n",
        "    labels = torch.concat(labels)\n",
        "  else:\n",
        "    preds, labels = preds[0], labels[0]\n",
        "  return preds, labels\n",
        "\n",
        "\n",
        "def export_adapt_head(smodel, submod):\n",
        "  w = submod.util.activate(\n",
        "    layer_names= [\"query\", 'value'],\n",
        "    key = 'export',\n",
        "    Modular=Lora)\n",
        "  if smodel.head is not None:\n",
        "    head_w = copy.deepcopy(smodel.head.state_dict())\n",
        "  else:\n",
        "    head_w  = None\n",
        "  return w, head_w\n",
        "\n",
        "def load_adapt_head(smodel,submod, adapt, head_w):\n",
        "  for w in adapt:\n",
        "    submod.get_submodule(w[0]).load_state_dict(w[1],  strict=False)\n",
        "\n",
        "  if smodel.head is not None:\n",
        "    smodel.head.load_state_dict(head_w)\n",
        "\n",
        "# load_adapt_head(smodel, w1, head1_w)\n",
        "\n"
      ],
      "metadata": {
        "id": "2c-9nuuoFzg8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 🟡 Pipe"
      ],
      "metadata": {
        "id": "_us60uYu2hI6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Models ref (teacher) 🎩"
      ],
      "metadata": {
        "id": "5n03ggfcKxhN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModel\n",
        "\n",
        "from transformers import  BitsAndBytesConfig\n",
        "\n",
        "\n",
        "model_name = 'allenai/specter2_base'\n",
        "\n",
        "# model_name = 'sentence-transformers/paraphrase-MiniLM-L3-v2'\n",
        "# model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "# model_name = 'allenai/scibert_scivocab_uncased'\n",
        "\n",
        "model_ref_t = AutoModel.from_pretrained(model_name)\n",
        "tokenizer_t_ref = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "# model_ref_t.encoder.layer[0].attention.output.dense.weight\n",
        "\n",
        "\n",
        "import copy\n",
        "\n",
        "max_seq_length = 512\n",
        "\n",
        "tokenize_t = lambda train_text: tokenizer_t_ref(train_text,\n",
        "    padding=\"max_length\",\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors=\"pt\")\n",
        "\n",
        "\n",
        "tokenize_batch_t = lambda train_text: tokenizer_t_ref.batch_encode_plus(\n",
        "    train_text,\n",
        "    add_special_tokens=True,\n",
        "    padding='max_length',\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors='pt',\n",
        "    return_attention_mask=True\n",
        ")\n",
        "\n",
        "# base = \"/content/gdrive/MyDrive/Projetos/TCC\"\n",
        "# base_save = f\"{base}/export/diagnostic/{model_name}\"\n",
        "# create_repo_if_not_exist(base_save+\"/\")\n",
        "# file_path = f\"/content/content/slr_data_tcc/{base_dataset}\"\n",
        "\n",
        "\n",
        "##### modular\n",
        "Mod = ModularUtils(Modular=ModMix)\n",
        "Mod.freeze(module=model_ref_t)\n",
        "\n",
        "smodel_mod_t = SentenceEmbedding(\n",
        "    model_ref_t,\n",
        "    normalize=False,\n",
        "    with_head = False\n",
        "    )\n",
        "\n",
        "\n",
        "smodel_mod_t, model_name"
      ],
      "metadata": {
        "id": "SSbJgJt8KzHz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "5uN6E1F9jgqf"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "2iIwt5v0JsAm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### big model"
      ],
      "metadata": {
        "id": "_DBmVvxTqpRh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "temperature of softmax\n",
        "\n",
        "the magnitude of positive gradient is equal to the sum of negative gradients. The temperature controls the distribution of negative gradients.\n",
        "\n",
        "Smaller temperature tends to concentrate more on the nearest neighbours of the anchor point, which plays a role in controlling the hardness-aware sensitivity\n",
        "\n",
        "\n",
        "This can be explained as follows: when the temperature is small, the contrastive loss tends to separate the positive samples\n",
        "close to the anchor sample, which makes the local distribution be sparse. With all samples are trained, the embedding\n",
        "space tends to make the neighbour of each point be sparse, and the distribution tends to be more uniform\n",
        "\n",
        "We could see that the tolerance is positively related to the temperature τ . However, the tolerance can not directly reflect the feature quality. For\n",
        "example, when all the samples reside in a single point of the hypersphere, then the tolerance is maximized, while the\n",
        "feature quality is bad. The tolerance reflects the local density of semantically related samples. An ideal model should\n",
        "be both locally clustered and globally uniform."
      ],
      "metadata": {
        "id": "xhLEbZqClR9x"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### data"
      ],
      "metadata": {
        "id": "mW8pEXo_D5yS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "\n",
        "# Set seed for NumPy\n",
        "\n",
        "# seed = 21442\n",
        "\n",
        "np.random.seed(seed)\n",
        "\n",
        "# Set seed for PyTorch\n",
        "torch.manual_seed(seed)\n",
        "\n",
        "# Set seed for CUDA if available\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed(seed)\n",
        "\n",
        "\n",
        "config = {\n",
        "    \"total_batch\": 16,\n",
        "    \"accumulation_steps\": 1,\n",
        "    \"batch_size\":int(16/1),\n",
        "    \"batch_size_eval\":16,\n",
        "###\n",
        "    \"n_epocs\": 2,\n",
        "    \"n_pairs\" : 40,\n",
        "    \"lr\" : 9e-3,\n",
        "###\n",
        "    \"alpha\":4*(0.1),\n",
        "    \"r\":4,\n",
        "    # \"tau\" : 0.05,\n",
        "    \"tau\" : 0.1,\n",
        "###\n",
        "    \"la\" : 3,\n",
        "    \"rest\" : 12 - 3,\n",
        "           }\n",
        "\n",
        "######\n",
        "##########################################################################\n",
        "##### modular\n",
        "\n",
        "base_dataset1 = \"cohen/ADHD.csv\"\n",
        "file_path = f\"data/content/slr_data_tcc/{base_dataset1}\"\n",
        "\n",
        "df_train1, exs_test1 = data_import(file_path, n=8, samples_n=300, samples_p=100)\n",
        "train_dataloader1, test_dataloader_eval1,  train_dataloader_eval1 = loaders(df_train1, exs_test1, config)\n",
        "\n",
        "\n",
        "# base_dataset = \"synergy/Smid_2020.csv\"\n",
        "base_dataset2 = \"SWIFT systematic review data/BPA.csv\"\n",
        "file_path = f\"data/content/slr_data_tcc/{base_dataset2}\"\n",
        "\n",
        "df_train2, exs_test2 = data_import(file_path, n=8, samples_n=300, samples_p=100)\n",
        "train_dataloader2, test_dataloader_eval2,  train_dataloader_eval2 = loaders(df_train2, exs_test2, config)\n",
        "\n",
        "\n",
        "base_dataset3 = \"SWIFT systematic review data/Fluoride.csv\"\n",
        "file_path = f\"data/content/slr_data_tcc/{base_dataset3}\"\n",
        "\n",
        "df_train3, exs_test3 = data_import(file_path, n=8, samples_n=300, samples_p=100)\n",
        "train_dataloader3, test_dataloader_eval3,  train_dataloader_eval3 = loaders(df_train3, exs_test3, config)\n",
        "\n",
        "\n",
        "base_dataset4 = \"SWIFT systematic review data/PFOS-PFOA.csv\"\n",
        "file_path = f\"data/content/slr_data_tcc/{base_dataset4}\"\n",
        "\n",
        "df_train4, exs_test4 = data_import(file_path, n=8, samples_n=300, samples_p=100)\n",
        "train_dataloader4, test_dataloader_eval4,  train_dataloader_eval4 = loaders(df_train4, exs_test4, config)\n",
        "\n",
        "\n",
        "# base_dataset5 = \"cohen/Opiods.csv\"\n",
        "# file_path = f\"data/content/slr_data_tcc/{base_dataset5}\"\n",
        "\n",
        "\n",
        "# df_train5, exs_test5 = data_import(file_path, n=8, samples_n=300, samples_p=100)\n",
        "# train_dataloader5, test_dataloader_eval5,  train_dataloader_eval5 = loaders(df_train5, exs_test5, config)\n",
        "\n",
        "\n",
        "base_dataset6 = \"SWIFT systematic review data/Transgenerational.csv\"\n",
        "file_path = f\"data/content/slr_data_tcc/{base_dataset6}\"\n",
        "\n",
        "df_train6, exs_test6 = data_import(file_path, n=8, samples_n=300, samples_p=100)\n",
        "train_dataloader6, test_dataloader_eval6,  train_dataloader_eval6 = loaders(df_train6, exs_test6, config)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "lod1 = (train_dataloader_eval1, test_dataloader_eval1)\n",
        "lod_train1= (df_train1, train_dataloader1, *lod1)\n",
        "\n",
        "lod2 = (train_dataloader_eval2, test_dataloader_eval2)\n",
        "lod_train2= (df_train2, train_dataloader2, *lod2)\n",
        "\n",
        "lod3 = (train_dataloader_eval3, test_dataloader_eval3)\n",
        "lod_train3= (df_train3, train_dataloader3, *lod3)\n",
        "\n",
        "lod4 = (train_dataloader_eval4, test_dataloader_eval4)\n",
        "lod_train4= (df_train4, train_dataloader4, *lod4)\n",
        "\n",
        "# lod5 = (train_dataloader_eval5, test_dataloader_eval5)\n",
        "# lod_train5= (df_train5, train_dataloader5, *lod5)\n",
        "\n",
        "lod6 = (train_dataloader_eval6, test_dataloader_eval6)\n",
        "lod_train6= (df_train6, train_dataloader6, *lod6)"
      ],
      "metadata": {
        "id": "1hMblxzAExBi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### No train"
      ],
      "metadata": {
        "id": "k_wr704MxE1R"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "k-aCKwITKPzX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModel\n",
        "\n",
        "\n",
        "\n",
        "model_name = 'allenai/specter2_base'\n",
        "# model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "\n",
        "\n",
        "\n",
        "model_ref_t = AutoModel.from_pretrained(model_name)\n",
        "tokenizer_t_ref = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "config = {\n",
        "    \"total_batch\": 16,\n",
        "    \"accumulation_steps\": 1,\n",
        "    \"batch_size\":int(16/1),\n",
        "    \"batch_size_eval\":16,\n",
        "###\n",
        "    \"n_epocs\": 2,\n",
        "    \"n_pairs\" : 40,\n",
        "    \"lr\" : 9e-3,\n",
        "###\n",
        "    \"alpha\":4*(0.1),\n",
        "    \"r\":4,\n",
        "    # \"tau\" : 0.05,\n",
        "    \"tau\" : 0.2,\n",
        "###\n",
        "    \"la\" : 3,\n",
        "    \"rest\" : 12 - 3,\n",
        "           }\n",
        "\n",
        "\n",
        "\n",
        "import copy\n",
        "\n",
        "max_seq_length = 512\n",
        "\n",
        "tokenize_t = lambda train_text: tokenizer_t_ref(train_text,\n",
        "    padding=\"max_length\",\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors=\"pt\")\n",
        "\n",
        "\n",
        "tokenize_batch_t = lambda train_text: tokenizer_t_ref.batch_encode_plus(\n",
        "    train_text,\n",
        "    add_special_tokens=True,\n",
        "    padding='max_length',\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors='pt',\n",
        "    return_attention_mask=True\n",
        ")\n",
        "\n",
        "# base = \"/content/gdrive/MyDrive/Projetos/TCC\"\n",
        "# base_save = f\"{base}/export/diagnostic/{model_name}\"\n",
        "# create_repo_if_not_exist(base_save+\"/\")\n",
        "# file_path = f\"/content/content/slr_data_tcc/{base_dataset}\"\n",
        "\n",
        "\n",
        "##### modular\n",
        "Mod = ModularUtils(Modular=ModMix)\n",
        "Mod.freeze(module=model_ref_t)\n",
        "\n",
        "smodel = SentenceEmbedding(\n",
        "    copy.deepcopy(model_ref_t),\n",
        "    normalize=False,\n",
        "    with_head = False\n",
        "    )\n",
        "\n",
        "smodel = smodel.half()\n",
        "\n"
      ],
      "metadata": {
        "id": "2qJ9WCjA3Cfe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### speed"
      ],
      "metadata": {
        "id": "2i6XEtZ4wtxO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def data_import(file_path, n=8, samples_n=200, samples_p=1):\n",
        "    print(\"####\"*5)\n",
        "    print()\n",
        "    print(os.path.basename(file_path))\n",
        "    df = preprocess_csv_file(file_path)\n",
        "\n",
        "\n",
        "    tk = lambda x: len(tokenizer_t_ref.encode(x))\n",
        "    tk_len = df.text.map(tk)\n",
        "    tk_len.hist()\n",
        "\n",
        "\n",
        "    df_rest = df[tk_len > 20]\n",
        "    df_train = df_rest.groupby('label').sample(n)\n",
        "\n",
        "    print(\"data count:\")\n",
        "    print(df['label'].value_counts())\n",
        "\n",
        "    print(df_train.iloc[0])\n",
        "\n",
        "    df_test = df[~df.index.isin(df_train.index)]\n",
        "\n",
        "    if len(df_test[df_test['label'] == 1]) < samples_p:\n",
        "        df_test_one = df_test[df_test['label'] == 1].copy()\n",
        "    else:\n",
        "        df_test_one = df_test[df_test['label'] == 1].sample(n=samples_p, replace=False)\n",
        "\n",
        "\n",
        "    if len(df_test[df_test['label'] == 0]) < samples_n:\n",
        "        df_test_zero = df_test[df_test['label'] == 0].copy()\n",
        "    else:\n",
        "        df_test_zero = df_test[df_test['label'] == 0].sample(n=samples_n, replace=False)\n",
        "\n",
        "    exs_test = pd.concat([df_test_one,df_test_zero])\n",
        "\n",
        "    print(\"Eval count:\")\n",
        "    print(exs_test['label'].value_counts())\n",
        "    return df_train, exs_test\n",
        "\n",
        "\n",
        "def loaders(df_train, exs_test, config):\n",
        "    sim_data = SIMDataset(df_train, tokenize_batch_t, config[\"n_pairs\"],\n",
        "                          )\n",
        "    train_dataloader = DataLoader(sim_data,shuffle=True,\n",
        "                                  batch_size=config[\"batch_size\"])\n",
        "\n",
        "    test_dataloader_eval = DataLoader(shuffle=False, batch_size=config[\"batch_size_eval\"],\n",
        "                dataset=SIMDataset(\n",
        "                    exs_test, tokenize_batch_t, 0,\n",
        "                          )\n",
        "                )\n",
        "\n",
        "    train_dataloader_eval = DataLoader(shuffle=False, batch_size=config[\"batch_size\"],\n",
        "                dataset=SIMDataset(\n",
        "                    df_train, tokenize_batch_t, 0,\n",
        "                    # emb= emb, device=\"cuda\", pttn=\"* d a\"\n",
        "                          )\n",
        "                )\n",
        "\n",
        "    return train_dataloader, test_dataloader_eval,  train_dataloader_eval\n",
        "\n",
        "\n",
        "\n",
        "# path= \"data/content/slr_data_tcc/cohen/Opiods.csv\"\n",
        "path= \"./content/slr_data_tcc/cohen/Opiods.csv\"\n",
        "df = preprocess_csv_file(path)\n",
        "\n",
        "tk = lambda x: len(tokenizer_t_ref.encode(x))\n",
        "tk_len = df.text.map(tk)\n",
        "\n",
        "ex = df[tk_len>=512]\n",
        "\n",
        "ex = ex.iloc[:1,:]\n",
        "all = []\n",
        "for _ in range(200):\n",
        "    all.append(ex)\n",
        "ex = pd.concat(all)\n",
        "\n",
        "eval_speed_dataloader = DataLoader(SIMDataset(ex, tokenize_batch_t, 0),\n",
        "                              shuffle=False,\n",
        "                              batch_size=200)\n",
        "\n",
        "\n",
        "smodel.to(torch.bfloat16)\n",
        "\n",
        "import time\n",
        "\n",
        "start_time = time.time()\n",
        "\n",
        "\n",
        "for x in eval_speed_dataloader:\n",
        "  smodel(x[0])\n",
        "end_time = time.time()\n",
        "execution_time = end_time - start_time\n",
        "print(\"Execution time:\", execution_time, \"seconds\")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Ll0w6_bdKOie"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install onnx\n",
        "!pip install onnxruntime"
      ],
      "metadata": {
        "id": "PnfE-np_Ufe9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import BertModel, BertTokenizer\n",
        "\n",
        "input_text = \"Your input text with 512 tokens.\"\n",
        "\n",
        "smodel.to(torch.float32)\n",
        "\n",
        "\n",
        "# Tokenize input with truncation to max_length\n",
        "inputs = tokenizer_t_ref(input_text, truncation=True, return_tensors=\"pt\",  padding='max_length', max_length=512)\n",
        "\n",
        "\n",
        "# Export the model to ONNX format\n",
        "torch.onnx.export(smodel.model, (inputs[\"input_ids\"], inputs[\"attention_mask\"]), \"bert_model.onnx\",\n",
        "                  input_names = ['input_ids', \"attention_mask\"],\n",
        "                                    dynamic_axes={\"input_ids\": {0: \"batch_size\"}, \"attention_mask\": {0: \"batch_size\"}})\n",
        "\n",
        "\n",
        "# Quantize the ONNX model to 8-bit integers\n",
        "from onnxruntime.quantization import quantize_dynamic\n",
        "\n",
        "quantize_dynamic(\"bert_model.onnx\", \"quantized_bert_model.onnx\")\n"
      ],
      "metadata": {
        "id": "n_EaupKIN_be"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import onnxruntime\n",
        "import numpy as np\n",
        "import time\n",
        "\n",
        "# Load the quantized BERT model\n",
        "ort_session = onnxruntime.InferenceSession(\"quantized_bert_model.onnx\")\n",
        "\n",
        "# Generate fake input data with 512 tokens\n",
        "fake_data = \" \".join([\"token\" + str(i) for i in range(512)])\n",
        "\n",
        "# Tokenize input data\n",
        "input_single = tokenizer_t_ref(fake_data, return_tensors=\"np\",max_length=512)\n",
        "\n",
        "# Repeat the single example to create a batch of 32 examples\n",
        "input_ids = np.repeat(input_single[\"input_ids\"], 32, axis=0)\n",
        "attention_mask = np.repeat(input_single[\"attention_mask\"], 32, axis=0)\n",
        "\n",
        "input_ids = input_ids.reshape(32, 512)\n",
        "attention_mask = attention_mask.reshape(32, 512)\n",
        "\n",
        "# Run inference and measure time\n",
        "start_time = time.time()\n",
        "outputs = ort_session.run(None, {\"input_ids\": input_ids.astype(np.int64), \"attention_mask\": attention_mask.astype(np.int64)})\n",
        "end_time = time.time()\n",
        "\n",
        "# Calculate inference time\n",
        "inference_time = end_time - start_time\n",
        "print(\"Total inference time for a batch of 32 examples:\", inference_time, \"seconds\")\n"
      ],
      "metadata": {
        "id": "k8varcFnOj8i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# Export the model to ONNX format\n",
        "torch.onnx.export(smodel.model, (inputs[\"input_ids\"], inputs[\"attention_mask\"]), \"bert_model.onnx\", verbose=True)\n",
        "\n",
        "# Load the ONNX model\n",
        "import onnx\n",
        "\n",
        "onnx_model = onnx.load(\"bert_model.onnx\")\n",
        "\n",
        "# Quantize the ONNX model to 8-bit integers\n",
        "from onnxruntime.quantization import quantize_dynamic\n",
        "\n",
        "quantized_model = quantize_dynamic(onnx_model, per_channel=True)\n",
        "\n",
        "# Save the quantized model\n",
        "onnx.save(quantized_model, \"quantized_bert_model.onnx\")\n",
        "\n",
        "# Now you can use the quantized_bert_model.onnx for inference on 8-bit CPUs\n"
      ],
      "metadata": {
        "id": "mmPRmjDpNsLt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### normal"
      ],
      "metadata": {
        "id": "P59_7o3rwxxo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os\n",
        "from pprint import pprint\n",
        "\n",
        "# base_path = 'data/content/slr_data_tcc'\n",
        "base_path = './content/slr_data_tcc'\n",
        "\n",
        "# Paths and names of the domains datasets\n",
        "# path = 'SLR_data'\n",
        "paths_cohen = glob.glob(f\"{base_path}/cohen/*.csv\", recursive=True)\n",
        "# pprint(paths_cohen)\n",
        "\n",
        "paths_swift = glob.glob(f\"{base_path}/SWIFT systematic review data/*.csv\", recursive=True)\n",
        "# pprint(paths_swift)\n",
        "\n",
        "\n",
        "# names = [os.path.basename(p) for p in paths]\n",
        "\n",
        "import random\n",
        "paths = [*paths_swift, *paths_cohen]\n",
        "\n",
        "paths_n = random.sample(paths, k=5)\n",
        "\n",
        "pprint(paths_n)"
      ],
      "metadata": {
        "id": "6JmvwBSnzK62"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# df_train, exs_test = data_import(p, n=8, samples_n=300, samples_p=100)\n",
        "\n",
        "\n",
        "\n",
        "def visualize_tsne(representations, labels,\n",
        "                   n_components=2,\n",
        "                   perplexity=30, save=None, alpha=1):\n",
        "    tsne = TSNE(n_components=n_components, perplexity=perplexity, random_state=42)\n",
        "    embeddings_2d = tsne.fit_transform(representations)\n",
        "\n",
        "    plt.figure(figsize=(10, 10))\n",
        "\n",
        "    for i in labels.unique():\n",
        "        indices = labels == i\n",
        "        plt.scatter(embeddings_2d[indices, 0],\n",
        "                    embeddings_2d[indices, 1], label=f'Classe {int(i)}', alpha=alpha)\n",
        "\n",
        "    plt.yticks(fontsize=12)\n",
        "    plt.xticks(fontsize=12)\n",
        "\n",
        "    # plt.title(f't-SNE do Model Embeddings ({n_components} components)')\n",
        "    legend = plt.legend(fontsize=14)\n",
        "\n",
        "    for text in legend.get_texts():\n",
        "      text.set_fontweight('bold')\n",
        "\n",
        "\n",
        "    if save is not None:\n",
        "      plt.savefig(f'{save}/tsne.pdf')\n",
        "\n",
        "    plt.show()"
      ],
      "metadata": {
        "id": "O1ikieIIJ5L6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "lods=[]\n",
        "\n",
        "for p in paths_n:\n",
        "\n",
        "  # p, l = eval_model_emb(smodel_test, [lod1,lod2])\n",
        "\n",
        "  df_train, exs_test = data_import(p, n=8, samples_n=300, samples_p=100)\n",
        "  train_dataloader, test_dataloader_eval,  train_dataloader_eval = loaders(df_train, exs_test, config)\n",
        "\n",
        "  lod = (train_dataloader_eval, test_dataloader_eval)\n",
        "\n",
        "  lods.append(lod)\n",
        "\n",
        "\n",
        "p, l = eval_model_emb(smodel, lods)\n",
        "\n",
        "# save = f\"plots/results/emd_wo_train\"\n",
        "# create_repo_if_not_exist(save)\n",
        "visualize_tsne(p, l, perplexity =7)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# smodel_test = SentenceEmbedding(\n",
        "#     copy.deepcopy(model_ref_t),\n",
        "#     normalize=False,\n",
        "#     with_head = False,\n",
        "#     )\n",
        "\n",
        "\n",
        "# Mod(smodel_test)\n",
        "# smodel_test.util.map(\n",
        "#     name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "# # activate in the last layers\n",
        "# submod_test = smodel_test.model.encoder.layer[-config[\"la\"]:]\n",
        "# Mod(submod_test)\n",
        "# submod_test.util.activate(\n",
        "#     layer_names= [\"query\", 'value'],\n",
        "#     key = 'set',-\n",
        "#     config = {\n",
        "#         'name':'test',\n",
        "#         \"alpha\":16*(1),\n",
        "#         \"r\":16,\n",
        "#         })\n",
        "# load_adapt_head(smodel_test, submod_test, *w_5)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "JJ4TDhny1yBH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "# path= \"data/content/slr_data_tcc/cohen/Opiods.csv\"\n",
        "path= \"./content/slr_data_tcc/cohen/NSAIDS.csv\"\n",
        "\n",
        "df_train, exs_test = data_import(path, n=8, samples_n=300, samples_p=100)\n",
        "train_dataloader, test_dataloader_eval,  train_dataloader_eval = loaders(df_train, exs_test, config)\n",
        "\n",
        "lod_eval = (train_dataloader_eval, test_dataloader_eval)\n"
      ],
      "metadata": {
        "id": "7iTxcUl4KWet"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def visualize_tsne(representations, labels,\n",
        "                   n_components=2,\n",
        "                   perplexity=30, save=None, alpha=1):\n",
        "    tsne = TSNE(n_components=n_components, perplexity=perplexity, random_state=42)\n",
        "    embeddings_2d = tsne.fit_transform(representations)\n",
        "\n",
        "    plt.figure(figsize=(10, 10))\n",
        "\n",
        "    for i in labels.unique():\n",
        "        indices = labels == i\n",
        "        plt.scatter(embeddings_2d[indices, 0],\n",
        "                    embeddings_2d[indices, 1],\n",
        "                    label=f'Classe {int(i)}', alpha=alpha, s = 100)\n",
        "\n",
        "    plt.yticks(fontsize=18)\n",
        "    plt.xticks(fontsize=18)\n",
        "\n",
        "    # plt.title(f't-SNE do Model Embeddings ({n_components} components)')\n",
        "    legend = plt.legend(fontsize=18)\n",
        "\n",
        "    for text in legend.get_texts():\n",
        "      text.set_fontweight('bold')\n",
        "\n",
        "\n",
        "    if save is not None:\n",
        "      plt.savefig(f'{save}/tsne.pdf')\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test = smodel_test.to(torch.float32)\n",
        "smodel_test = smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.model.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "# submod =submod.half()\n",
        "\n",
        "p0, l0 = eval_model_emb(smodel, [lod_eval])\n",
        "# visualize_tsne(p, l, perplexity =20, save=None)\n",
        "\n",
        "save = f\"plots/results/emd_example/data_wo/\"\n",
        "create_repo_if_not_exist(save)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "lod_train= (df_train, train_dataloader, train_dataloader_eval, test_dataloader_eval)\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train ,config)\n",
        "\n",
        "\n",
        "p1, l1 = eval_model_emb(smodel_test, [lod_eval])\n",
        "\n",
        "save = f\"plots/results/emd_example/data_w/\"\n",
        "create_repo_if_not_exist(save)\n",
        "\n",
        "visualize_tsne(p0, l0, perplexity =15, save=save)\n",
        "visualize_tsne(p1, l1, perplexity =15, save=save)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "5J5-JzVI4RJA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 12*((64**2+64**2+64**2+64**2)*12 + 768*3072 + 3070*768)/10**8"
      ],
      "metadata": {
        "id": "wtwlSsG7gnpM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = smodel_test\n",
        "submod_test = model.model.encoder.layer[-config[\"la\"]:]\n",
        "lod = lod_eval\n",
        "\n",
        "\n",
        "# w = merge([\n",
        "#                    w_1[0],\n",
        "#                    w_2[0],\n",
        "#                   #  w_3[0],\n",
        "#                   #  w_4[0],\n",
        "#                   #  w_6[0],\n",
        "#     ],\n",
        "#                   [\n",
        "#                    w_1[1],\n",
        "#                    w_2[1],\n",
        "#                   #  w_3[1],\n",
        "#                   #  w_4[1],\n",
        "#                   #  w_6[1]\n",
        "#                       ],\n",
        "#                   [\n",
        "#                    0.7,\n",
        "#                    0.7,\n",
        "#                   #  0.2,\n",
        "#                   # 0.2,\n",
        "#                   # 0.2\n",
        "#                   #  1\n",
        "#                   ],\n",
        "#                   r=32, p=0.5\n",
        "#                   )\n",
        "\n",
        "# load_adapt_head(model, submod_test, *w)\n",
        "\n",
        "\n",
        "from pprint import pprint\n",
        "_=model.cuda()\n",
        "res = eval_pipeline(model,*lod,\n",
        "                    device=\"cuda\", save=None)\n",
        "pprint(res)\n",
        "_=model.cpu()\n",
        "\n",
        "\n",
        "\n",
        "model.eval()\n",
        "model.cuda()\n",
        "emb_test, emb_train = evaluate_model(model ,*lod, \"cuda\")\n",
        "model.cpu()\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "ids_pos = sum(emb_test[1]==1)\n",
        "print(f\"fisst positive samples: 0:{ids_pos}\")\n",
        "\n",
        "# Example usage:\n",
        "# Assuming emb_test[0] contains the batch of vectors\n",
        "batch_vectors_normalized = F.normalize(emb_test[0], p=2, dim=1)\n",
        "similarities = torch.matmul(batch_vectors_normalized, batch_vectors_normalized.t())\n",
        "plot_cosine_similarity_matrix(similarities[:,:], -1, 1)\n"
      ],
      "metadata": {
        "id": "hlJbS72JE9d2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = smodel\n",
        "submod_test = model.model.encoder.layer[-config[\"la\"]:]\n",
        "lod = lod_eval\n",
        "\n",
        "\n",
        "# w = merge([\n",
        "#                    w_1[0],\n",
        "#                    w_2[0],\n",
        "#                   #  w_3[0],\n",
        "#                   #  w_4[0],\n",
        "#                   #  w_6[0],\n",
        "#     ],\n",
        "#                   [\n",
        "#                    w_1[1],\n",
        "#                    w_2[1],\n",
        "#                   #  w_3[1],\n",
        "#                   #  w_4[1],\n",
        "#                   #  w_6[1]\n",
        "#                       ],\n",
        "#                   [\n",
        "#                    0.7,\n",
        "#                    0.7,\n",
        "#                   #  0.2,\n",
        "#                   # 0.2,\n",
        "#                   # 0.2\n",
        "#                   #  1\n",
        "#                   ],\n",
        "#                   r=32, p=0.5\n",
        "#                   )\n",
        "\n",
        "# load_adapt_head(model, submod_test, *w)\n",
        "\n",
        "\n",
        "from pprint import pprint\n",
        "_=model.cuda()\n",
        "res = eval_pipeline(model,*lod,\n",
        "                    device=\"cuda\", save=None)\n",
        "pprint(res)\n",
        "_=model.cpu()\n",
        "\n",
        "\n",
        "\n",
        "model.eval()\n",
        "model.cuda()\n",
        "emb_test, emb_train = evaluate_model(model ,*lod, \"cuda\")\n",
        "model.cpu()\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "ids_pos = sum(emb_test[1]==1)\n",
        "print(f\"fisst positive samples: 0:{ids_pos}\")\n",
        "\n",
        "# Example usage:\n",
        "# Assuming emb_test[0] contains the batch of vectors\n",
        "batch_vectors_normalized = F.normalize(emb_test[0], p=2, dim=1)\n",
        "similarities = torch.matmul(batch_vectors_normalized,\n",
        "                            batch_vectors_normalized.t())\n",
        "plot_cosine_similarity_matrix(similarities[:,:], -1, 1)\n"
      ],
      "metadata": {
        "id": "kD6cutj-LLOZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### choose model"
      ],
      "metadata": {
        "id": "ElMEXgzTB1nJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModel\n",
        "from transformers import  BitsAndBytesConfig\n",
        "\n",
        "\n",
        "model_names = [\n",
        "    'allenai/specter2_base',\n",
        "    'sentence-transformers/paraphrase-MiniLM-L3-v2',\n",
        "    'sentence-transformers/all-MiniLM-L12-v2',\n",
        "    'allenai/scibert_scivocab_uncased'\n",
        "    ]\n",
        "\n",
        "\n",
        "for model_name in model_names:\n",
        "  model_ref_t = AutoModel.from_pretrained(model_name)\n",
        "  tokenizer_t_ref = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "  # model_ref_t.encoder.layer[0].attention.output.dense.weight\n",
        "\n",
        "\n",
        "  import copy\n",
        "\n",
        "  max_seq_length = 512\n",
        "\n",
        "  tokenize_t = lambda train_text: tokenizer_t_ref(train_text,\n",
        "      padding=\"max_length\",\n",
        "      max_length=max_seq_length,\n",
        "      truncation=True,\n",
        "      return_tensors=\"pt\")\n",
        "\n",
        "\n",
        "  tokenize_batch_t = lambda train_text: tokenizer_t_ref.batch_encode_plus(\n",
        "      train_text,\n",
        "      add_special_tokens=True,\n",
        "      padding='max_length',\n",
        "      max_length=max_seq_length,\n",
        "      truncation=True,\n",
        "      return_tensors='pt',\n",
        "      return_attention_mask=True\n",
        "  )\n",
        "\n",
        "  # base = \"/content/gdrive/MyDrive/Projetos/TCC\"\n",
        "  # base_save = f\"{base}/export/diagnostic/{model_name}\"\n",
        "  # create_repo_if_not_exist(base_save+\"/\")\n",
        "  # file_path = f\"/content/content/slr_data_tcc/{base_dataset}\"\n",
        "\n",
        "\n",
        "  ##### modular\n",
        "  Mod = ModularUtils(Modular=ModMix)\n",
        "  Mod.freeze(module=model_ref_t)\n",
        "\n",
        "  smodel_mod_t = SentenceEmbedding(\n",
        "      model_ref_t,\n",
        "      normalize=False,\n",
        "      with_head = True\n",
        "      )\n",
        "\n",
        "\n",
        "  smodel_mod_t, model_name\n",
        "\n",
        "  ###### model\n",
        "  smodel = copy.deepcopy(smodel_mod_t)\n",
        "  head = copy.deepcopy(smodel.head)\n",
        "  submod = smodel.model.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "\n",
        "  base_save = \"metrics\"\n",
        "  save = f\"{base_save}/{model_name}/\"\n",
        "  create_repo_if_not_exist(save)\n",
        "\n",
        "\n",
        "  lod= (df_train5, train_dataloader5, train_dataloader_eval5, test_dataloader_eval5)\n",
        "  res = main_pipe_train(smodel, submod, *lod ,config, save)\n",
        "  # w_4 = export_adapt_head(smodel, submod)\n",
        "  export_json(res[0], f\"{save}Metrics.json\")"
      ],
      "metadata": {
        "id": "8aXNFCP0B5Qf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os\n",
        "from pprint import pprint\n",
        "\n",
        "base_path = 'metrics/**/Metrics.json'\n",
        "\n",
        "# Paths and names of the domains datasets\n",
        "# path = 'SLR_data'\n",
        "paths_metrics = glob.glob(base_path, recursive=True)\n",
        "pprint(paths_metrics)\n",
        "\n",
        "metrics_list=[]\n",
        "names=[]\n",
        "for p in paths_metrics:\n",
        "  metric = read_json(p)\n",
        "  name = p.split(\"/\")[-2]\n",
        "  metrics_list.append(metric[\"observed\"])\n",
        "  names.append(name)\n",
        "\n",
        "def plot_rocs(x_list, names):\n",
        "    plt.figure(figsize=(7, 5))\n",
        "    for x, name in zip(x_list, names):\n",
        "      fpr = x.get(\"fpr\")\n",
        "      tpr = x.get(\"tpr\")\n",
        "      roc_auc = x.get(\"roc_auc\")\n",
        "      plt.plot(fpr, tpr, lw=2, label=f'{name} (roc_auc = {roc_auc:.2f})')\n",
        "\n",
        "    plt.axhline(y=0.95, color='red', linestyle='--', label=f'Recall 95%')\n",
        "    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')\n",
        "\n",
        "    plt.xlabel('Falso Positivo (%)', fontsize=14)\n",
        "    plt.ylabel('Verdadeiro Positivo (%)', fontsize=14)\n",
        "    plt.yticks(fontsize=12)\n",
        "    plt.xticks(fontsize=12)\n",
        "\n",
        "    plt.legend(loc='lower right')\n",
        "    plt.tight_layout()\n",
        "    plt.savefig('model_choose.pdf')\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "plot_rocs(metrics_list,\n",
        "          names)\n",
        "\n"
      ],
      "metadata": {
        "id": "C8VkjG0nJNtK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 16bit"
      ],
      "metadata": {
        "id": "vkuFWKgp2f-d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModel\n",
        "\n",
        "from transformers import  BitsAndBytesConfig\n",
        "\n",
        "\n",
        "model_name = 'allenai/specter2_base'\n",
        "# model_name = 'sentence-transformers/paraphrase-MiniLM-L3-v2'\n",
        "# model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "# model_name = 'allenai/scibert_scivocab_uncased'\n",
        "\n",
        "\n",
        "\n",
        "config = {\n",
        "    \"total_batch\": 16,\n",
        "    \"accumulation_steps\": 1,\n",
        "    \"batch_size\":int(16/1),\n",
        "    \"batch_size_eval\":16,\n",
        "###\n",
        "    \"n_epocs\": 2,\n",
        "    \"n_pairs\" : 40,\n",
        "    \"lr\" : 9e-3,\n",
        "###\n",
        "    \"alpha\":4*(0.1),\n",
        "    \"r\":4,\n",
        "    \"tau\" : 0.1,\n",
        "###\n",
        "    \"la\" : 3,\n",
        "    \"rest\" : 12 - 3,\n",
        "           }\n",
        "\n",
        "model_ref_t = AutoModel.from_pretrained(model_name)\n",
        "tokenizer_t_ref = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "import copy\n",
        "\n",
        "max_seq_length = 512\n",
        "\n",
        "tokenize_t = lambda train_text: tokenizer_t_ref(train_text,\n",
        "    padding=\"max_length\",\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors=\"pt\")\n",
        "\n",
        "\n",
        "tokenize_batch_t = lambda train_text: tokenizer_t_ref.batch_encode_plus(\n",
        "    train_text,\n",
        "    add_special_tokens=True,\n",
        "    padding='max_length',\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors='pt',\n",
        "    return_attention_mask=True\n",
        ")\n",
        "\n",
        "##### modular\n",
        "Mod = ModularUtils(Modular=ModMix)\n",
        "Mod.freeze(module=model_ref_t)\n",
        "\n",
        "\n",
        "smodel = SentenceEmbedding(\n",
        "    model_ref_t,\n",
        "    normalize=False,\n",
        "    with_head = False\n",
        "    )\n",
        "\n",
        "\n",
        "smodel, model_name"
      ],
      "metadata": {
        "id": "LWnd5tN_RUvq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# config[\"lr\"]\n",
        "\n",
        "\n",
        "# base_save = \"metrics\"\n",
        "# # base_dataset1\n",
        "# save = f\"{base_save}/{model_name}/\"\n",
        "# create_repo_if_not_exist(save)\n",
        "# export_json(res, f\"{save}Metrics.json\")\n",
        "\n",
        "\n",
        "###### model\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test=smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train1 ,config)\n",
        "w_1 = export_adapt_head(smodel_test, submod)\n",
        "\n",
        "##### model\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test=smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train2 ,config)\n",
        "w_2 = export_adapt_head(smodel_test, submod)\n",
        "\n",
        "\n",
        "###### model\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test=smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train3 ,config)\n",
        "w_3 = export_adapt_head(smodel_test, submod)\n",
        "\n",
        "\n",
        "###### model\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test=smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train4 ,config)\n",
        "w_4 = export_adapt_head(smodel_test, submod)\n",
        "\n",
        "\n",
        "# ###### model\n",
        "# smodel_test = copy.deepcopy(smodel)\n",
        "# smodel_test=smodel_test.to(torch.bfloat16)\n",
        "# submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "# _ = main_pipe_train(smodel_test, submod, *lod_train5 ,config)\n",
        "# w_5 = export_adapt_head(smodel_test, submod)\n",
        "\n",
        "\n",
        "###### model\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test=smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train6 ,config)\n",
        "w_6 = export_adapt_head(smodel_test, submod)\n"
      ],
      "metadata": {
        "id": "3GhyHFa08es1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "def zero_out_smallest_percentage(tensor, percentage):\n",
        "    \"\"\"\n",
        "    Zero out a certain percentage of smallest values in a PyTorch tensor.\n",
        "\n",
        "    Args:\n",
        "    - tensor: input PyTorch tensor\n",
        "    - percentage: percentage of smallest values to zero out\n",
        "\n",
        "    Returns:\n",
        "    - masked_tensor: tensor with specified percentage of smallest values zeroed out\n",
        "    \"\"\"\n",
        "    # Calculate the number of elements to zero out\n",
        "    num_zeros = int(tensor.numel() * percentage)\n",
        "\n",
        "    # Find the indices of the smallest values\n",
        "    _, indices = torch.topk(tensor.abs().reshape(-1), num_zeros, largest=False)\n",
        "\n",
        "    # Create a mask of zeros and ones with the same shape as the tensor\n",
        "    mask = torch.ones_like(tensor)\n",
        "\n",
        "    # Set selected indices to 0\n",
        "    mask.reshape(-1)[indices] = 0\n",
        "\n",
        "    # Apply the mask to zero out selected values\n",
        "    masked_tensor = tensor * mask\n",
        "\n",
        "    return masked_tensor\n",
        "\n",
        "# Example usage:\n"
      ],
      "metadata": {
        "id": "1ixoxztkH-oX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from einop import einop\n",
        "def merge(ws, heads, weights=[], r=8, p=0):\n",
        "  w = copy.deepcopy(ws[0])\n",
        "  head = copy.deepcopy(heads[0])\n",
        "\n",
        "  if len(weights)==0:\n",
        "    weights=[1/len(ws)]*len(ws)\n",
        "\n",
        "  for i in range(len(w)):\n",
        "    a = torch.stack([w_temp[i][1][\"loraA\"]* weights[j] for j, w_temp in enumerate(ws)])\n",
        "    b = torch.stack([w_temp[i][1][\"loraB\"]* weights[j] for j, w_temp in enumerate(ws)])\n",
        "\n",
        "    # a = einop(a, \"b i r -> i (b r)\")\n",
        "    # b = einop(b, \"b r o ->(b r) o\")\n",
        "\n",
        "\n",
        "    dtype= a.dtype\n",
        "    a= a.to(torch.float32)\n",
        "    b= b.to(torch.float32)\n",
        "    # a = torch.quantile(a, p, dim=0)\n",
        "    # b = torch.quantile(b, p, dim=0)\n",
        "    # a[torch.abs(a)<torch.quantile(torch.abs(a), 0.2)] = 0\n",
        "    # b[torch.abs(b)<torch.quantile(torch.abs(b), 0.2)] = 0\n",
        "\n",
        "\n",
        "    ww = a@b\n",
        "    ww = torch.quantile(ww, p, dim=0)\n",
        "    # ww = zero_out_smallest_percentage(ww, p)\n",
        "    ww[torch.abs(ww)<torch.quantile(torch.abs(ww), 0.5)] = 0\n",
        "\n",
        "\n",
        "    U,S,V = torch.svd(ww)\n",
        "    a =U[:, :r]\n",
        "    b=  torch.diag(S[:r]) @ V[:, :r].T\n",
        "\n",
        "\n",
        "    a= a.to(dtype)\n",
        "    b= b.to(dtype)\n",
        "\n",
        "    w[i][1][\"loraA\"] = a\n",
        "    w[i][1][\"loraB\"] = b\n",
        "\n",
        "  return w, head\n",
        "\n",
        "\n",
        "# head_w_opi = dict(head_w_opi)\n",
        "\n",
        "r=8\n",
        "w = merge([\n",
        "                   w_1[0],\n",
        "                   w_2[0],\n",
        "                   w_3[0],\n",
        "                   w_4[0],\n",
        "                  #  w_5[0],\n",
        "                   w_6[0],\n",
        "    ],\n",
        "                  [\n",
        "                   w_1[1],\n",
        "                   w_2[1],\n",
        "                   w_3[1],\n",
        "                   w_4[1],\n",
        "                  #  w_5[1],\n",
        "                   w_6[1]\n",
        "                      ],\n",
        "                  [\n",
        "                   1,\n",
        "                   1,\n",
        "                   1,\n",
        "                   1,\n",
        "                   1,\n",
        "                   1,\n",
        "                  ],\n",
        "                  r=r, p=0.7\n",
        "                  )\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "smodel_test_emb = SentenceEmbedding(\n",
        "    copy.deepcopy(model_ref_t),\n",
        "    with_head = False,\n",
        "    normalize=False,\n",
        "    )\n",
        "\n",
        "\n",
        "\n",
        "Mod(smodel_test_emb)\n",
        "smodel_test_emb.util.map(\n",
        "    name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "# activate in the last layers\n",
        "submod_test = smodel_test_emb.model.encoder.layer[-config[\"la\"]:]\n",
        "Mod(submod_test)\n",
        "submod_test.util.activate(\n",
        "    layer_names= [\"query\", 'value'],\n",
        "    key = 'set',\n",
        "    config = {\n",
        "        'name':'test',\n",
        "        \"r\":r,\n",
        "        \"alpha\":r*(0.1)\n",
        "        })\n",
        "\n",
        "\n",
        "load_adapt_head(smodel_test_emb, submod_test, *w)\n",
        "\n",
        "########################## eval\n",
        "\n",
        "smodel=smodel.to(torch.float16)\n",
        "smodel_test_emb=smodel_test_emb.to(torch.float16)\n",
        "p0, l0 = eval_model_emb(smodel, [lod1,lod2,lod3,lod4,lod6])\n",
        "p1, l1 = eval_model_emb(smodel_test_emb, [lod1,lod2,lod3,lod4,lod6])\n",
        "from sklearn.manifold import TSNE\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "\n",
        "def visualize_tsne(representations, labels,\n",
        "                   n_components=2,\n",
        "                   perplexity=30, save=None, alpha=1):\n",
        "    tsne = TSNE(n_components=n_components, perplexity=perplexity, random_state=42)\n",
        "    embeddings_2d = tsne.fit_transform(representations)\n",
        "\n",
        "    plt.figure(figsize=(10, 10))\n",
        "\n",
        "    for i in [-1,1,-2,2,-3,3,-4,4,-5,5]:\n",
        "        indices = labels == i\n",
        "        plt.scatter(embeddings_2d[indices, 0],\n",
        "                    embeddings_2d[indices, 1],\n",
        "                    label=f'Classe {int(i)}', alpha=alpha, s=100)\n",
        "\n",
        "    plt.yticks(fontsize=18)\n",
        "    plt.xticks(fontsize=18)\n",
        "    # plt.xlim(-50, 50)\n",
        "    # plt.ylim(-50, 50)\n",
        "\n",
        "    # Place legend outside of plot\n",
        "    # legend = plt.legend(bbox_to_anchor=(0.5, 1.2), loc='upper center', fontsize=24, ncol=5)\n",
        "\n",
        "    # for text in legend.get_texts():\n",
        "    #   text.set_fontweight('bold')\n",
        "\n",
        "    if save is not None:\n",
        "      plt.savefig(f'{save}/tsne.pdf')\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "\n",
        "visualize_tsne(p0, l0, perplexity =40, save=None, alpha=0.7)\n",
        "visualize_tsne(p1, l1, perplexity =40, save=None, alpha=0.7)"
      ],
      "metadata": {
        "id": "cD5NBvUFhDy5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "smodel_test_emb = smodel_test_emb.to(torch.float16)\n",
        "model = smodel_test_emb\n",
        "submod_test = model.model.encoder.layer[-config[\"la\"]:]\n",
        "lod = lod1\n",
        "\n",
        "\n",
        "load_adapt_head(model, submod_test, *w)\n",
        "\n",
        "\n",
        "# from pprint import pprint\n",
        "# _=model.cuda()\n",
        "# res = eval_pipeline(model,*lod,\n",
        "#                     device=\"cuda\", save=None)\n",
        "# pprint(res)\n",
        "# _=model.cpu()\n",
        "\n",
        "\n",
        "\n",
        "model.eval()\n",
        "model.cuda()\n",
        "emb_test, emb_train = evaluate_model(model ,*lod, \"cuda\")\n",
        "model.cpu()\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "ids_pos = sum(emb_test[1]==1)\n",
        "print(f\"fisst positive samples: 0:{ids_pos}\")\n",
        "\n",
        "# Example usage:\n",
        "# Assuming emb_test[0] contains the batch of vectors\n",
        "batch_vectors_normalized = F.normalize(emb_test[0], p=2, dim=1)\n",
        "similarities = torch.matmul(batch_vectors_normalized, batch_vectors_normalized.t())\n",
        "plot_cosine_similarity_matrix(similarities[:,:], -1, 1)\n"
      ],
      "metadata": {
        "id": "Swhq5XXDTsFt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "smodel_test_emb = SentenceEmbedding(\n",
        "    copy.deepcopy(model_ref_t),\n",
        "    with_head = False,\n",
        "    normalize=False,\n",
        "    )\n",
        "\n",
        "\n",
        "\n",
        "Mod(smodel_test_emb)\n",
        "smodel_test_emb.util.map(\n",
        "    name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "# activate in the last layers\n",
        "submod_test = smodel_test_emb.model.encoder.layer[-config[\"la\"]:]\n",
        "Mod(submod_test)\n",
        "submod_test.util.activate(\n",
        "    layer_names= [\"query\", 'value'],\n",
        "    key = 'set',\n",
        "    config = {\n",
        "        'name':'test',\n",
        "        \"r\":4,\n",
        "        \"alpha\":4*(0.1)\n",
        "        })\n",
        "\n",
        "\n",
        "\n",
        "lod = lod1\n",
        "smodel_test_emb = smodel_test_emb.to(torch.float16)\n",
        "model = smodel_test_emb\n",
        "\n",
        "\n",
        "load_adapt_head(model, submod_test, *w_1)\n",
        "\n",
        "\n",
        "# from pprint import pprint\n",
        "# _=model.cuda()\n",
        "# res = eval_pipeline(model,*lod,\n",
        "#                     device=\"cuda\", save=None)\n",
        "# pprint(res)\n",
        "# _=model.cpu()\n",
        "\n",
        "\n",
        "\n",
        "model.eval()\n",
        "model.cuda()\n",
        "emb_test, emb_train = evaluate_model(model ,*lod, \"cuda\")\n",
        "model.cpu()\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "ids_pos = sum(emb_test[1]==1)\n",
        "print(f\"fisst positive samples: 0:{ids_pos}\")\n",
        "\n",
        "# Example usage:\n",
        "# Assuming emb_test[0] contains the batch of vectors\n",
        "batch_vectors_normalized = F.normalize(emb_test[0], p=2, dim=1)\n",
        "similarities = torch.matmul(batch_vectors_normalized, batch_vectors_normalized.t())\n",
        "plot_cosine_similarity_matrix(similarities[:,:])\n"
      ],
      "metadata": {
        "id": "f_2lRIOR0l13"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from pprint import pprint\n",
        "smodel_test_emb = SentenceEmbedding(\n",
        "    copy.deepcopy(model_ref_t),\n",
        "    with_head = False,\n",
        "    normalize=False,\n",
        "    )\n",
        "\n",
        "smodel_test_emb = smodel_test_emb.to(torch.float16)\n",
        "\n",
        "\n",
        "model = smodel_test_emb\n",
        "\n",
        "\n",
        "res_cru = []\n",
        "\n",
        "for lo, wo in zip([lod1, lod2, lod3, lod4, lod6],\n",
        "                  [w_1, w_2, w_3, w_4, w_6]):\n",
        "  _=model.cuda()\n",
        "  res = eval_pipeline(model,*lo,\n",
        "                      device=\"cuda\", save=None)\n",
        "  pprint(res)\n",
        "  _=model.cpu()\n",
        "  res_cru.append(res)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "###########################\n",
        "\n",
        "\n",
        "smodel_test_emb = SentenceEmbedding(\n",
        "    copy.deepcopy(model_ref_t),\n",
        "    with_head = False,\n",
        "    normalize=False,\n",
        "    )\n",
        "\n",
        "smodel_test_emb = smodel_test_emb.to(torch.float16)\n",
        "\n",
        "\n",
        "Mod(smodel_test_emb)\n",
        "smodel_test_emb.util.map(\n",
        "    name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "# activate in the last layers\n",
        "submod_test = smodel_test_emb.model.encoder.layer[-config[\"la\"]:]\n",
        "Mod(submod_test)\n",
        "submod_test.util.activate(\n",
        "    layer_names= [\"query\", 'value'],\n",
        "    key = 'set',\n",
        "    config = {\n",
        "        'name':'test',\n",
        "        \"r\":4,\n",
        "        \"alpha\":4*(0.1)\n",
        "        })\n",
        "\n",
        "model = smodel_test_emb\n",
        "submod_test\n",
        "smodel_test_emb = smodel_test_emb.to(torch.float16)\n",
        "\n",
        "\n",
        "\n",
        "res_just_one = []\n",
        "\n",
        "for lo, wo in zip([lod1, lod2, lod3, lod4, lod6],\n",
        "                  [w_1, w_2, w_3, w_4, w_6]):\n",
        "  load_adapt_head(model, submod_test, *wo)\n",
        "  _=model.cuda()\n",
        "  res = eval_pipeline(model,*lo,\n",
        "                      device=\"cuda\", save=None)\n",
        "  pprint(res)\n",
        "  _=model.cpu()\n",
        "  res_just_one.append(res)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "###########################\n",
        "\n",
        "smodel_test_emb = SentenceEmbedding(\n",
        "    copy.deepcopy(model_ref_t),\n",
        "    with_head = False,\n",
        "    normalize=False,\n",
        "    )\n",
        "\n",
        "\n",
        "\n",
        "Mod(smodel_test_emb)\n",
        "smodel_test_emb.util.map(\n",
        "    name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "# activate in the last layers\n",
        "submod_test = smodel_test_emb.model.encoder.layer[-config[\"la\"]:]\n",
        "Mod(submod_test)\n",
        "submod_test.util.activate(\n",
        "    layer_names= [\"query\", 'value'],\n",
        "    key = 'set',\n",
        "    config = {\n",
        "        'name':'test',\n",
        "        \"r\":r,\n",
        "        \"alpha\":r*(0.1)\n",
        "        })\n",
        "\n",
        "\n",
        "\n",
        "smodel_test_emb = smodel_test_emb.to(torch.float16)\n",
        "model = smodel_test_emb\n",
        "\n",
        "\n",
        "load_adapt_head(model, submod_test, *w)\n",
        "\n",
        "\n",
        "res_merged = []\n",
        "for lo, wo in zip([lod1, lod2, lod3, lod4, lod6],\n",
        "                  [w_1, w_2, w_3, w_4, w_6]):\n",
        "  _=model.cuda()\n",
        "  res = eval_pipeline(model,*lo,\n",
        "                      device=\"cuda\", save=None)\n",
        "  pprint(res)\n",
        "  _=model.cpu()\n",
        "  res_merged.append(res)\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Bdk1iKjTg96k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "names = [base_dataset1,base_dataset2,base_dataset3,base_dataset4,base_dataset6]\n",
        "\n",
        "names_d = [os.path.basename(x).split(\".\")[0] for x in names]\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "QxBYaqa9zC6w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def plot_lines_with_labels(data_lists, label_list):\n",
        "    fig, ax = plt.subplots(figsize=(12, 8))\n",
        "\n",
        "    line_styles = ['-', ':', '--']  # Define line styles for each line\n",
        "    # colors = ['blue', 'green', 'orange', 'purple']  # Define colors for each line\n",
        "    # colors = plt.cm.viridis(np.linspace(0, 1, 7))\n",
        "    colors = ['green', 'red','blue', 'purple', 'brown']\n",
        "\n",
        "    legds=[]\n",
        "    for j, data_list in enumerate(data_lists):\n",
        "        color = colors[j % len(colors)]  # Cycle through colors\n",
        "\n",
        "        auc_values = []  # Store AUC values\n",
        "        for i, data in enumerate(data_list):\n",
        "            fpr = data.get(\"fpr\")\n",
        "            tpr = data.get(\"tpr\")\n",
        "            roc_auc = data.get(\"roc_auc\")\n",
        "            line_style = line_styles[i % len(line_styles)]  # Cycle through line styles\n",
        "            ax.plot(fpr, tpr, lw=2, alpha=0.7,\n",
        "                    linestyle=line_style, color=color)\n",
        "            auc_values.append(f'{roc_auc:.2f}')  # Append AUC value to list\n",
        "\n",
        "        # Format legend entry with all AUC values\n",
        "        legend_label = f'{label_list[j]} (AUCs: {\" / \".join(auc_values)})'\n",
        "        ls, = ax.plot([], [], lw=2, linestyle='-', color=color, label=legend_label)\n",
        "        legds.append(ls)\n",
        "\n",
        "    recall=ax.axhline(y=0.95, color='black', linestyle='-', label=f'Recall 95%', lw=1)\n",
        "\n",
        "    legend1 = ax.legend(handles=[*legds, recall],\n",
        "                        loc='lower right', fontsize=14)  # Adjust location and ncols\n",
        "\n",
        "\n",
        "    ax.plot([0, 1], [0, 1], color='navy', lw=1, linestyle='--')\n",
        "\n",
        "    line1, = ax.plot([], [], lw=2, linestyle=line_styles[0], color=\"black\", label=\"Base\")\n",
        "    line2, = ax.plot([], [], lw=2, linestyle=line_styles[1], color=\"black\", label=\"Lora\")\n",
        "    line3, = ax.plot([], [], lw=2, linestyle=line_styles[2], color=\"black\", label=\"Híbrido\")\n",
        "\n",
        "    # Create a legend for the first line.\n",
        "    legend2 = ax.legend(handles=[line1, line2, line3],\n",
        "                        loc='upper center', bbox_to_anchor=(0.5, -0.15), ncol=3, fontsize=14)  # Adjust location and ncols\n",
        "\n",
        "\n",
        "\n",
        "    ax.add_artist(legend1)\n",
        "    # ax.add_artist(legend2)\n",
        "\n",
        "    # Add the legend manually to the Axes.\n",
        "\n",
        "    ax.set_xlabel('Falso Positivo (%)', fontsize=16, weight=\"bold\")\n",
        "    ax.set_ylabel('Verdadeiro Positivo (%)', fontsize=16, weight=\"bold\")\n",
        "    plt.yticks(fontsize=18)\n",
        "    plt.xticks(fontsize=18)\n",
        "\n",
        "\n",
        "# Test the function\n",
        "i=0\n",
        "data_lists = [\n",
        "    [res_cru[i][\"observed\"], res_just_one[i][\"observed\"], res_merged[i][\"observed\"]],\n",
        "    [res_cru[i+1][\"observed\"], res_just_one[i+1][\"observed\"], res_merged[i+1][\"observed\"]],\n",
        "    [res_cru[i+2][\"observed\"], res_just_one[i+2][\"observed\"], res_merged[i+2][\"observed\"]],\n",
        "    [res_cru[i+3][\"observed\"], res_just_one[i+3][\"observed\"], res_merged[i+3][\"observed\"]],\n",
        "    [res_cru[i+4][\"observed\"], res_just_one[i+4][\"observed\"], res_merged[i+4][\"observed\"]],\n",
        "    # Add more data lists as needed\n",
        "]\n",
        "\n",
        "label_list = [\n",
        "    f\"{names_d[i]}\",\n",
        "    f\"{names_d[i+1]}\",\n",
        "    f\"{names_d[i+2]}\",\n",
        "    f\"{names_d[i+3]}\",\n",
        "    f\"{names_d[i+4]}\",\n",
        "    ]  # Provide labels for each data set\n",
        "\n",
        "plot_lines_with_labels(data_lists, label_list)\n",
        "# '-', '--', '-.', ':', 'None', ' ', '', 'solid', 'dashed', 'dashdot', 'dotted'\n"
      ],
      "metadata": {
        "id": "kwocdtrIKd2U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dfs = []\n",
        "\n",
        "\n",
        "for i in range(5):\n",
        "  df = pd.DataFrame({\"awss\": [res_cru[i][\"95\"][\"awss\"]],\n",
        "                     \"wss\": [res_cru[i][\"95\"][\"wss\"]],\n",
        "                    \"dataset\": [names_d[i]],\n",
        "                     \"método\":\"base\"}\n",
        "                    )\n",
        "\n",
        "  dfs.append(df)\n",
        "\n",
        "for i in range(5):\n",
        "  df = pd.DataFrame({\"awss\": [res_just_one[i][\"95\"][\"awss\"]],\n",
        "                     \"wss\": [res_just_one[i][\"95\"][\"wss\"]],\n",
        "                    \"dataset\": [names_d[i]],\n",
        "                     \"método\":\"Lora\"}\n",
        "                    )\n",
        "\n",
        "  dfs.append(df)\n",
        "for i in range(5):\n",
        "  df = pd.DataFrame({\"awss\": [res_merged[i][\"95\"][\"awss\"]],\n",
        "                     \"wss\": [res_merged[i][\"95\"][\"wss\"]],\n",
        "                    \"dataset\": [names_d[i]],\n",
        "                     \"método\":\"Híbrido\"}\n",
        "                    )\n",
        "\n",
        "  dfs.append(df)\n",
        "\n",
        "\n",
        "\n",
        "result_df = pd.concat(dfs, ignore_index=True)\n",
        "# print(result_df)\n",
        "\n",
        "result_df=result_df.pivot(index='dataset', columns=['método'], values=['awss'])\n",
        "print(result_df)\n",
        "print(result_df.to_latex(float_format=\"%.2f\"))\n",
        "\n",
        "\n",
        "# # Pivot the DataFrame\n",
        "# pivot_df = result_df.pivot_table(index='dataset', columns='método', values=['awss', 'wss'])\n",
        "\n",
        "\n",
        "\n",
        "# print(pivot_df.to_latex())\n",
        "\n",
        "\n",
        "# # Rearrange columns\n",
        "# pivot_df = pivot_df.swaplevel(axis=1).sort_index(axis=1)\n",
        "# pivot_df"
      ],
      "metadata": {
        "id": "OAYRO9MxU_wV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "res_cru[i][\"95\"][\"awss\"]\n",
        "\n",
        "\n",
        "\n",
        "res_just_one[i][\"95\"]\n",
        "res_merged[i][\"95\"]"
      ],
      "metadata": {
        "id": "Ewkk7vcWUVHB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "#### 🔱 optimization\n"
      ],
      "metadata": {
        "id": "0-Io8O8IzzxH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### quantization"
      ],
      "metadata": {
        "id": "gljQcB8F2ksn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# model_base_t = copy.deepcopy(model_ref_t)\n",
        "\n",
        "\n",
        "##########################################################################\n",
        "# ## Quant\n",
        "smodel_mod_t.util.activate(\n",
        "        layer_names= [\"query\", 'value']+[\"key\", 'dense'],\n",
        "        key = 'set',\n",
        "        config = {\n",
        "            'name': 'Quant',\n",
        "            'bits': 4,\n",
        "            \"copyw\":True,\n",
        "            \"base\":\"base_linear\"\n",
        "            })\n",
        "\n",
        "smodel_mod_t.util.activate(\n",
        "        layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "        key = 'quantize',\n",
        "        Modular=Quantize_Linear)\n",
        "\n",
        "##########################################################################\n",
        "\n",
        "## lora\n",
        "submod.util.activate(\n",
        "        layer_names= [\"query\", 'value'],\n",
        "        key = 'set', config= {'name': 'test',\n",
        "                              'base': 'Quant'}\n",
        ")\n",
        "\n",
        "\n",
        "##########################################################################\n",
        "smodel_mod_t.cuda()\n",
        "\n",
        "smodel_mod_t.eval()\n",
        "res = eval_pipeline(smodel_mod_t,\n",
        "                    train_dataloader=train_dataloader_eval,\n",
        "                    test_dataloader=test_dataloader_eval,\n",
        "              device=\"cuda\")\n",
        "\n",
        "from pprint import pprint\n",
        "pprint(res)\n",
        "\n",
        "smodel_mod_t.to(\"cpu\")\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "# export_json(res, \"export/diagnostic/Qlora.json\")"
      ],
      "metadata": {
        "id": "Q_QGHMK1LvNH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# m = copy.deepcopy(smodel_mod_t)\n",
        "\n",
        "# # del m\n",
        "# m.util.activate(\n",
        "#         layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "#         key = 'export',\n",
        "#         Modular=QLinear)\n",
        "\n",
        "\n",
        "# m.util.activate(\n",
        "#         layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "#         key = 'del_layer',\n",
        "#         config={\n",
        "#             \"\",\n",
        "#             \"base\":\"Quant\"\n",
        "#         })\n",
        "\n",
        "# m\n",
        "\n",
        "# # torch.save(m.state_dict(), \"model.bin\")"
      ],
      "metadata": {
        "id": "ajIxo-zI4Sg5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### sparse"
      ],
      "metadata": {
        "id": "8zPmIPWu2wDF"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "uPbXhPCGeq-k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "smodel_mod_t.util.activate(\n",
        "        layer_names= [\"query\", 'value', \"key\"],\n",
        "        key = 'set',\n",
        "        config = {\n",
        "            'name': 'Prune',\n",
        "            's': 0.7,\n",
        "            'acum': True,\n",
        "            \"copyw\":True\n",
        "            })\n",
        "\n",
        "\n",
        "smodel_mod_t.util.activate(\n",
        "        layer_names= ['dense'],\n",
        "        key = 'set',\n",
        "        config = {\n",
        "            'name': 'Prune',\n",
        "            's': 0.5,\n",
        "            'acum': True,\n",
        "            \"copyw\":True\n",
        "            })\n",
        "\n",
        "smodel_mod_t.to(\"cuda\")\n",
        "get_batch_labels(\n",
        "    smodel_mod_t,\n",
        "    train_dataloader_eval)\n",
        "smodel_mod_t.to(\"cpu\")\n",
        "\n",
        "\n",
        "smodel_mod_t.util.activate(\n",
        "        layer_names= ['dense']+[\"query\", 'value', \"key\"],\n",
        "        key = 'prune',\n",
        "        Modular=Prune\n",
        "        )\n",
        "\n",
        "smodel_mod_t.util.activate(\n",
        "        layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "        key = 'set',\n",
        "        config = {\n",
        "            'name': 'Prune',\n",
        "            'fake_prune': True,\n",
        "            })\n",
        "\n",
        "## lora\n",
        "submod.util.activate(\n",
        "        layer_names= [\"query\", 'value'],\n",
        "        key = 'set', config= {'name': 'test',\n",
        "                              'base': 'Prune'\n",
        "                              }\n",
        ")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "smodel_mod_t.cuda()\n",
        "smodel_mod_t.eval()\n",
        "res = eval_pipeline(smodel_mod_t,\n",
        "                    train_dataloader=train_dataloader_eval,\n",
        "                    test_dataloader=test_dataloader_eval,\n",
        "              device=\"cuda\")\n",
        "smodel_mod_t.to(\"cpu\")\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "########\n",
        "m = copy.deepcopy(smodel_mod_t)\n",
        "\n",
        "m.util.activate(\n",
        "        layer_names= ['dense']+[\"query\", 'value', \"key\"],\n",
        "        key = 'prune',\n",
        "        Modular=Prune\n",
        "        )\n",
        "m.model.encoder.layer[0].attention.self.query.Mix.Prune.main.weight\n",
        "\n",
        "sub_copy = m.model.encoder.layer[-3:]\n",
        "\n",
        "## lora\n",
        "Mod(sub_copy)\n",
        "m.util.activate(\n",
        "        layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "        key = 'del_layer',\n",
        "        config={\n",
        "            \"name\":[\"base_linear\", \"Quant\"],\n",
        "            \"base\":\"Prune\"\n",
        "        })\n",
        "\n",
        "sub_copy.util.activate(\n",
        "        layer_names= [\"query\", 'value'],\n",
        "        key = 'set', config= {'name': 'test',\n",
        "                              'base': 'Prune'}\n",
        ")\n",
        "\n",
        "sub_copy.util.activate(\n",
        "        layer_names= [\"query\", 'value'],\n",
        "        key = 'del_layer',\n",
        "        config={\n",
        "            \"name\":[\"Adapt\"],\n",
        "            \"base\":\"Prune\"\n",
        "        })\n",
        "##########\n",
        "\n",
        "\n",
        "res[\"sparse\"]= 100 - count_non_zero_params(m)\n",
        "del m\n",
        "\n",
        "from pprint import pprint\n",
        "pprint(res)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "CoOx6VwQ0ifv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### 🎒 Distillation"
      ],
      "metadata": {
        "id": "hGuVrjxlEeJ3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "[La mini paper](https://arxiv.org/pdf/2002.10957.pdf) + [Sbert](https://arxiv.org/pdf/1908.10084.pdf)\n",
        "\n",
        "finetuned in 1B pairs\n"
      ],
      "metadata": {
        "id": "bsJD9danINm1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "###### pipe"
      ],
      "metadata": {
        "id": "_Fjh3J1SrSfE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tk = lambda x: len(tokenizer_t.encode(x))\n",
        "\n",
        "tk_len = df_test.text.map(tk)\n",
        "\n",
        "tk_len.hist()\n",
        "df_rest_unlb = df_test[tk_len > 100]"
      ],
      "metadata": {
        "id": "608Ahg0hX3zL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "df_train_unlabel = df_rest_unlb.sample(64)\n"
      ],
      "metadata": {
        "id": "XCczx_mMVFf8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.optim import Adam\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "\n",
        "model_kind = \"SIMModel\"\n",
        "\n",
        "\n",
        "total_batch = 32\n",
        "accumulation_steps = 1\n",
        "\n",
        "batch_size=int(total_batch/accumulation_steps)\n",
        "\n",
        "\n",
        "######\n",
        "n_epocs= 2\n",
        "n_pairs = 30\n",
        "n_unlabaled = 8\n",
        "\n",
        "lr = 3e-3\n",
        "\n",
        "\n",
        "###\n",
        "lora_alpha=1\n",
        "r=8\n",
        "tau = 0.1\n",
        "\n",
        "###\n",
        "la = 3\n",
        "rest = 12 - la\n",
        "\n",
        "######\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# pd.concat([df_test.sample(n_unlabaled),df_train])\n",
        "# just labeled data\n",
        "n_epocs= 2\n",
        "sim_data_1 = SIMDataset(df_train, tokenize_batch_t, n_pairs)\n",
        "train_dataloader_1 = DataLoader(sim_data_1, shuffle=True, batch_size=batch_size)\n",
        "\n",
        "interation_1 = int(len(sim_data_1)/batch_size)\n",
        "total_iterations_1 = n_epocs*(interation_1)\n",
        "\n",
        "\n",
        "# just unlabeled data\n",
        "\n",
        "n_epocs= 2\n",
        "sim_data_2 = SIMDataset(df_train_unlabel, tokenize_batch_t, n_pairs = 300, n_pairs_random = True)\n",
        "train_dataloader_2 = DataLoader(sim_data_2, shuffle=True, batch_size=batch_size)\n",
        "\n",
        "interation_2 = int(len(sim_data_2)/batch_size)\n",
        "total_iterations_2 = n_epocs*(interation_2)\n",
        "\n",
        "# unlabeled + labeld data\n",
        "accumulation_steps= 1\n",
        "batch_size=int(64/accumulation_steps)\n",
        "n_epocs=2\n",
        "sim_data_3 = SIMDataset(pd.concat([df_train_unlabel, df_train]),\n",
        "                        tokenize_batch_t,\n",
        "                        n_pairs = 600,\n",
        "                        n_pairs_random = True)\n",
        "train_dataloader_3 = DataLoader(sim_data_3, shuffle=True, batch_size=batch_size)\n",
        "\n",
        "interation_3 = int(len(sim_data_3)/batch_size)\n",
        "total_iterations_3 = n_epocs*(interation_3)\n",
        "\n"
      ],
      "metadata": {
        "id": "fnLI5MmnrkXp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Mod = ModularUtils()\n",
        "model_name = 'sentence-transformers/paraphrase-MiniLM-L3-v2'\n",
        "# model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "modelref_s = AutoModel.from_pretrained(model_name)\n",
        "tokenizer_s = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "Mod.freeze_layers(modelref_s, freeze=True)\n",
        "Mod.freeze_layers(smodel_mod_t, freeze=True)\n",
        "\n",
        "model_base_s = copy.deepcopy(modelref_s)\n",
        "model_mod_s = Mod.copy_arch(model_base_s)\n",
        "smodel_mod_s = SentenceEmbedding(\n",
        "    model_mod_s,\n",
        "    hidden =384,\n",
        "    with_head=True\n",
        "    )\n",
        "\n",
        "# proj = nn.Sequential(\n",
        "#       nn.Linear(384, 384),\n",
        "#       # nn.GELU(),\n",
        "#       # nn.Linear(384, 384)\n",
        "#     )\n",
        "\n",
        "# smodel_mod_s.head.add_module(\n",
        "#     'proj',\n",
        "#     proj\n",
        "# )\n",
        "\n",
        "# model_base_t\n",
        "# model_mod_t\n",
        "# smodel_mod_t\n",
        "\n",
        "\n",
        "max_seq_length = 512\n",
        "\n",
        "tokenize_s = lambda train_text: tokenizer_s(train_text,\n",
        "    padding=\"max_length\",\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors=\"pt\")\n",
        "\n",
        "\n",
        "tokenize_batch_s = lambda train_text: tokenizer_s.batch_encode_plus(\n",
        "    train_text,\n",
        "    add_special_tokens=True,\n",
        "    padding='max_length',\n",
        "    max_length=max_seq_length,\n",
        "    truncation=True,\n",
        "    return_tensors='pt',\n",
        "    return_attention_mask=True\n",
        ")\n",
        "\n",
        "###### model\n",
        "\n",
        "\n",
        "Mod = ModularUtils()\n",
        "\n",
        "# base model and the lora adapter model\n",
        "Mod.freeze_layers(model_ref_t, freeze=True)\n",
        "\n",
        "\n",
        "model_base_s = copy.deepcopy(model_ref_s)\n",
        "model_mod_s = Mod.copy_arch(model_base_s)\n",
        "\n",
        "\n",
        "Mod(smodel_mod_s)\n",
        "\n",
        "\n",
        "# apply ModMix in all corresponding values\n",
        "\n",
        "smodel_mod_s.util.map(\n",
        "    Modular=ModMix,\n",
        "    name_module = [\"query\", 'value'])\n",
        "\n",
        "\n",
        "\n",
        "# Active Lora just in some layers\n",
        "\n",
        "# get the last layers and apply the lora\n",
        "submod = smodel_mod_s.get_submodule('encoder.layer')[-la:]\n",
        "\n",
        "\n",
        "# Mod(submod)\n",
        "submod.util.activate(\n",
        "    layer_names= [\"query\", 'value'],\n",
        "    key = 'set',\n",
        "    name= 'test', config = {\n",
        "    'alpha': 1,\n",
        "    'r':32})\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "print_trainable_params(smodel_mod_s)\n",
        "\n",
        "# l2 = loss_distill_to_light(\n",
        "#     model_s = smodel_mod_s,\n",
        "#     model_t = smodel_mod_t\n",
        "\n",
        "# )\n",
        "\n",
        "\n",
        "\n",
        "# with just df_train\n",
        "trainer_1 = Trainer(losses=few_shot_loss_distill(\n",
        "    model_s = smodel_mod_s,\n",
        "    model_t = smodel_mod_t,\n",
        "    loss=1).to('cuda'),\n",
        "                  optimizer=Adam(smodel_mod_s.parameters(), lr=lr),\n",
        "                  accumulation_steps = accumulation_steps,\n",
        "                  train_dataloader=train_dataloader_1,\n",
        "                  total_iterations=total_iterations_1,\n",
        "                  device = 'cuda')\n",
        "\n",
        "# with just df_unlabel\n",
        "# trainer_2 = Trainer(losses=few_shot_loss_distill(\n",
        "#     model_s = smodel_mod_s,\n",
        "#     model_t = smodel_mod_t,\n",
        "#     loss=1).to('cuda'),\n",
        "#                   optimizer=Adam(smodel_mod_s.parameters(), lr=lr),\n",
        "#                   accumulation_steps = accumulation_steps,\n",
        "#                   train_dataloader=train_dataloader_2,\n",
        "#                   total_iterations=total_iterations_2,\n",
        "#                   device = 'cuda')\n",
        "\n",
        "# with  dftrain and df_unlabel\n",
        "# 5e-4\n",
        "trainer_3 = Trainer(losses=few_shot_loss_distill(\n",
        "    model_s = smodel_mod_s,\n",
        "    model_t = smodel_mod_t,\n",
        "    loss=0).to('cuda'),\n",
        "                  optimizer=Adam(smodel_mod_s.parameters(), lr=lr),\n",
        "                  accumulation_steps = accumulation_steps,\n",
        "                  train_dataloader=train_dataloader_3,\n",
        "                  total_iterations=total_iterations_3,\n",
        "                  device = 'cuda')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#### train\n",
        "\n",
        "print(f\"Quantidade de amostras (rotuladas)\\n {len(sim_data_1)}\")\n",
        "print(f\"Quantidade de exemplos (rotuladas)\\n {len(df_train)}\")\n",
        "print(f\"Quantidade de amostras (nao rotuladas-par)\\n {len(sim_data_2)}\")\n",
        "print(f\"Quantidade de exemplos (nao rotuladas)\\n {len(df_train_unlabel)}\")\n",
        "print(f\"Quantidade de amostras (não + rotuladas-par)\\n {len(sim_data_3)}\")\n",
        "# print(f\"Quantidade de exemplos (nao rotuladas)\\n {len(df_train_unlabel)}\")\n",
        "\n",
        "\n",
        "\n",
        "print(submod[0].attention.self.value.loraA.weight)\n",
        "\n",
        "# model_s2.to(\"cpu\")\n",
        "# torch.cuda.empty_cache()\n",
        "# trainer_2.train()\n",
        "# trainer_1.train()\n",
        "\n",
        "# trainer_3.train()\n",
        "\n",
        "# smodel_mod_s.to('cuda')\n",
        "# res = eval_pipeline(smodel_mod_s,\n",
        "#               (df_train,pd.concat([df_train,exs_test])),\n",
        "#               is_SLRC=(model_kind == \"SLRC\"),\n",
        "#               batch_size=300,\n",
        "#                     tokenize = tokenize_batch_s)\n",
        "\n",
        "# from pprint import pprint\n",
        "# pprint(res)\n",
        "\n",
        "# smodel_mod_s.to(\"cpu\")\n",
        "# torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "smodel_mod_s.to('cuda')\n",
        "trainer_1.train()\n",
        "print(submod[0].attention.self.value.loraA.weight)\n",
        "\n",
        "res = eval_pipeline(smodel_mod_s,\n",
        "              (df_train,pd.concat([df_train,exs_test])),\n",
        "              is_SLRC=(model_kind == \"SLRC\"),\n",
        "              batch_size=300,\n",
        "                    tokenize = tokenize_batch_s)\n",
        "\n",
        "from pprint import pprint\n",
        "pprint(res)\n",
        "\n",
        "smodel_mod_s.to(\"cpu\")\n",
        "torch.cuda.empty_cache()\n"
      ],
      "metadata": {
        "id": "YUrE8oAVPFfa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for n, mo in list(submod.named_children()):\n",
        "  mo.attention.self.query.merge()\n",
        "  mo.attention.self.value.merge()\n",
        "\n",
        "\n",
        "proj2 = nn.Sequential(\n",
        "      nn.Linear(384, 384),\n",
        "      # nn.GELU(),\n",
        "      # nn.Linear(384, 384)\n",
        "    )\n",
        "\n",
        "Mod.freeze_layers(smodel_mod_s, freeze=True)\n",
        "smodel_mod_s.model = model_base_s\n",
        "\n",
        "\n",
        "submod2,_ =  Mod.app_Modular(\n",
        "        submod, # Modular\n",
        "        [\"base_linear\"], # linear layes (type_layers)\n",
        "        )\n",
        "\n",
        "\n",
        "\n",
        "# del smodel_mod_s.head.proj\n",
        "\n",
        "smodel_mod_s.head.add_module(\n",
        "    'proj',\n",
        "    proj2\n",
        ")\n",
        "\n",
        "\n",
        "l3 = loss_sim(smodel_mod_s,\n",
        "              tau = 0.05)\n",
        "print_trainable_params(l3)\n",
        "\n",
        "print(submod[0].attention.self.value.loraA.weight)\n",
        "\n",
        "trainer_4 = Trainer(losses=l3.to('cuda'),\n",
        "                  optimizer=Adam(l3.parameters(), lr=3e-5),\n",
        "                  accumulation_steps = accumulation_steps,\n",
        "                  train_dataloader=train_dataloader_1,\n",
        "                  total_iterations=total_iterations_1,\n",
        "                  device = 'cuda')\n",
        "trainer_4.train()\n",
        "\n",
        "\n",
        "print(submod[0].attention.self.value.loraA.weight)\n",
        "\n",
        "\n",
        "smodel_mod_s.to('cuda')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "res = eval_pipeline(smodel_mod_s,\n",
        "              (df_train,pd.concat([df_train,exs_test])),\n",
        "              is_SLRC=(model_kind == \"SLRC\"),\n",
        "              batch_size=300,\n",
        "                    tokenize = tokenize_batch_s)\n",
        "\n",
        "from pprint import pprint\n",
        "pprint(res)\n",
        "\n",
        "smodel_mod_s.to(\"cpu\")\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "Qae8BMmizeLu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://github.com/intel/neural-compressor?tab=readme-ov-file\n",
        "\n"
      ],
      "metadata": {
        "id": "Xwbuk8ozwbb6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inferece"
      ],
      "metadata": {
        "id": "hDqjb6E0EyX2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install optimum[onnxruntime-gpu]"
      ],
      "metadata": {
        "id": "ROHWe_9xxHjA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import onnxruntime as rt\n",
        "\n",
        "# #define the priority order for the execution providers\n",
        "# # prefer CUDA Execution Provider over CPU Execution Provider\n",
        "# EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']\n",
        "\n",
        "# # initialize the model.onnx\n",
        "# sess = rt.InferenceSession(\"model.onnx\", providers=EP_list)\n",
        "\n",
        "# # get the outputs metadata as a list of :class:`onnxruntime.NodeArg`\n",
        "# output_name = sess.get_outputs()[0].name\n",
        "\n",
        "# # get the inputs metadata as a list of :class:`onnxruntime.NodeArg`\n",
        "# input_name = sess.get_inputs()[0].name\n",
        "\n",
        "# # inference run using image_data as the input to the model\n",
        "# detections = sess.run([output_name], {input_name: image_data})[0]\n",
        "\n",
        "# print(\"Output shape:\", detections.shape)"
      ],
      "metadata": {
        "id": "b30SxKZzEz5S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 🌐 Main results"
      ],
      "metadata": {
        "id": "iBdeOEdspqbO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "- [x] escolher o modelo dentro de 1 dataset\n",
        "- [x] usar modelo dentro de todos os datasets/ avaliar metodos de otimização e metricas\n",
        "- [x] salvar loras e avaliar performance geral se fosse apenas 1 lora\n",
        "\n",
        "- [ ] avaliar aprendizado continuo em apenas um dataset\n",
        "- [ ] ?avaliar modelo destilado\n",
        "\n",
        "- [ ] falar sobre estimativa dos arquivos, quantidade de uso de memoria e velocidades e possibilidades de uso e deploy como modelo geral de transformaçãp"
      ],
      "metadata": {
        "id": "KXshDVUNErb6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## dataset statistics"
      ],
      "metadata": {
        "id": "wzKZLaERrJ5g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import os\n",
        "\n",
        "# def list_files_in_folders(base_path):\n",
        "#     folder_dict = {}\n",
        "\n",
        "#     def list_files(directory):\n",
        "#         return [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]\n",
        "\n",
        "#     def list_directories(directory):\n",
        "#         return [d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))]\n",
        "\n",
        "#     for folder in list_directories(base_path):\n",
        "#         folder_path = os.path.join(base_path, folder)\n",
        "#         files_in_folder = list_files(folder_path)\n",
        "#         folder_dict[folder] = files_in_folder\n",
        "\n",
        "#     return folder_dict\n",
        "\n",
        "# # Substitua '/content/content/slr_data_tcc' pelo seu caminho real\n",
        "# result_dict = list_files_in_folders('/content/content/slr_data_tcc')\n",
        "\n",
        "# # Imprima o resultado\n",
        "# for folder, files in result_dict.items():\n",
        "#     print(f\"{folder}: {files}\")\n"
      ],
      "metadata": {
        "id": "vRy1le_tnwk7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os"
      ],
      "metadata": {
        "id": "ozfSJ3MQsEdn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os\n",
        "from pprint import pprint\n",
        "\n",
        "# base_path = 'data/content/slr_data_tcc'\n",
        "base_path = './content/slr_data_tcc'\n",
        "\n",
        "# Paths and names of the domains datasets\n",
        "# path = 'SLR_data'\n",
        "paths_cohen = glob.glob(f\"{base_path}/cohen/*.csv\", recursive=True)\n",
        "pprint(paths_cohen)\n",
        "\n",
        "paths_swift = glob.glob(f\"{base_path}/SWIFT systematic review data/*.csv\", recursive=True)\n",
        "pprint(paths_swift)\n",
        "\n",
        "\n",
        "# names = [os.path.basename(p) for p in paths]\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "nfPeWer_sXMO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "names_r = {\n",
        "    \"Opiods\": \"Opioids\",\n",
        "    \"ACEInhibitors\": \"ACE Inhibitors\",\n",
        "    \"AtypicalAntipsychotics\": \"Atypical Antipsychotics\",\n",
        "    \"BetaBlockers\": \"Beta Blockers\",\n",
        "    \"UrinaryIncontinence\": \"Urinary Incontinence\",\n",
        "    \"OralHypoglycemics\": \"Oral Hypoglycemics\",\n",
        "    \"ProtonPumpInhibitors\": \"Proton Pump Inhibitors\",\n",
        "    \"CalciumChannelBlockers\": \"Calcium Channel Blockers\",\n",
        "    \"SkeletalMuscleRelaxants\": \"Skeletal Muscle Relaxants\",\n",
        "    \"SkeletalMuscleRelaxants\": \"Skeletal Muscle Relaxants\"\n",
        "    }"
      ],
      "metadata": {
        "id": "5t3wxnxttj_Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "tk = lambda x: len(tokenizer_t_ref.encode(x))\n",
        "\n",
        "names = []\n",
        "tokens_len_0 = []\n",
        "tokens_len_1 = []\n",
        "\n",
        "for p in paths_cohen:\n",
        "    n = os.path.basename(p).split(\".\")[0]\n",
        "    if n in list(names_r.keys()):\n",
        "      names.append(names_r[n])\n",
        "    else:\n",
        "      names.append(n)\n",
        "    df = preprocess_csv_file(p)\n",
        "    tokens_len_0.append(df[df[\"label\"] == 0].text.map(tk))\n",
        "    tokens_len_1.append(df[df[\"label\"] == 1].text.map(tk))\n",
        "\n",
        "for p in paths_swift:\n",
        "    n = os.path.basename(p).split(\".\")[0]\n",
        "    if n in list(names_r.keys()):\n",
        "      names.append(names_r[n])\n",
        "    else:\n",
        "      names.append(n)\n",
        "    df = preprocess_csv_file(p)\n",
        "    tokens_len_0.append(df[df[\"label\"] == 0].text.map(tk))\n",
        "    tokens_len_1.append(df[df[\"label\"] == 1].text.map(tk))\n"
      ],
      "metadata": {
        "id": "hPrcqH5KSON_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.patches as mpatches\n",
        "\n",
        "plt.figure(figsize=(10, 6))\n",
        "\n",
        "box_width = 0.4\n",
        "\n",
        "box_0 = plt.boxplot(tokens_len_0, positions=[pos - box_width / 2 for pos in range(len(names))],\n",
        "                    widths=box_width, patch_artist=True, boxprops=dict(facecolor='lightcoral'))\n",
        "box_1 = plt.boxplot(tokens_len_1, positions=[pos + box_width / 2 for pos in range(len(names))],\n",
        "                    widths=box_width, patch_artist=True, boxprops=dict(facecolor='skyblue'))\n",
        "\n",
        "# Customizing colors for the median line\n",
        "for median in [box_0['medians'], box_1['medians']]:\n",
        "    for median_line in median:\n",
        "        median_line.set(color='black', linewidth=2)\n",
        "\n",
        "plt.xticks([pos for pos in range(len(names))], [f'{name}' for name in names],\n",
        "           rotation=45, ha='right', fontsize=14, weight='bold')\n",
        "\n",
        "# Adjusting horizontal padding of tick labels\n",
        "plt.xticks(ha='right', rotation_mode='anchor', fontsize=14, weight='bold')\n",
        "plt.yticks(fontsize=12)\n",
        "\n",
        "plt.ylabel('Quant. de Tokens', fontsize=14, weight='bold')\n",
        "# plt.xlabel('Banco de Dados', fontsize=14, weight='bold')\n",
        "\n",
        "# Creating legend patches\n",
        "legend_patches = [mpatches.Patch(facecolor='skyblue', label='Label 1'),\n",
        "                  mpatches.Patch(facecolor='lightcoral', label='Label -1')]\n",
        "plt.legend(handles=legend_patches, loc='upper left', fontsize='large')\n",
        "\n",
        "plt.grid(True)\n",
        "\n",
        "plt.tight_layout()\n",
        "\n",
        "# Save the plot as PDF\n",
        "plt.savefig('boxplot_tokens.pdf')\n",
        "\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "fCG2RnLJSXKW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Given list\n",
        "data_list = [('ACE Inhibitors', (2214, 167, 2047)),\n",
        " ('ADHD', (781, 80, 701)),\n",
        " ('Antihistamines', (277, 87, 190)),\n",
        " ('Atypical Antipsychotics', (999, 329, 670)),\n",
        " ('Beta Blockers', (1819, 266, 1553)),\n",
        " ('Calcium Channel Blockers', (1069, 246, 823)),\n",
        " ('Estrogens', (337, 77, 260)),\n",
        " ('NSAIDs', (348, 81, 267)),\n",
        " ('Opioids', (1717, 41, 1676)),\n",
        " ('Oral Hypoglycemics', (462, 134, 328)),\n",
        " ('Proton Pump Inhibitors', (1171, 220, 951)),\n",
        " ('Skeletal Muscle Relaxants', (1318, 26, 1292)),\n",
        " ('Statins', (2659, 150, 2509)),\n",
        " ('Triptans', (573, 200, 373)),\n",
        " ('Urinary Incontinence', (271, 65, 206)),\n",
        " ('Drug Reviews (COHEN et al., 2006)', (16015, 2169, 13846)),\n",
        " ('Bisphenol A (BPA)', (7093, 102, 6991)),\n",
        " ('Fluoride and Neurotoxicity', (3870, 49, 3821)),\n",
        " ('Neurophatic pain', (29202, 5009, 24193)),\n",
        " ('PFOA/PFOS', (5950, 95, 5855)),\n",
        " ('Transgenerational', (46147, 606, 45541)),\n",
        " ('SWIFT (??)', (92262, 5861, 86401))]\n",
        "\n",
        "# Compute proportion of positives for each entry\n",
        "proportions = []\n",
        "for entry in data_list:\n",
        "    total, positive, _ = entry[1]\n",
        "    proportion = positive / total\n",
        "    proportions.append(proportion)\n",
        "\n",
        "# Display the proportions\n",
        "for entry, proportion in zip(data_list, proportions):\n",
        "    print(f\"{entry[0]}: &{proportion:.4f}\")\n"
      ],
      "metadata": {
        "id": "1yZq6dAJZdGH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## avaliation all"
      ],
      "metadata": {
        "id": "afnh5NRykrJL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.optim import Adam\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from tqdm import tqdm\n",
        "from torch.optim.lr_scheduler import LambdaLR\n",
        "#######\n",
        "import pandas as pd\n",
        "import os\n",
        "from glob import glob\n",
        "from pprint import pprint\n",
        "\n",
        "def process_files(base_path):\n",
        "    paths = glob(f\"{base_path}/**/*.csv\", recursive=True)\n",
        "\n",
        "    # Criar listas para armazenar os DataFrames e os caminhos dos arquivos\n",
        "    dfs = []\n",
        "    file_paths = []\n",
        "\n",
        "    # Iterar sobre os caminhos dos arquivos\n",
        "    for file_path in paths:\n",
        "        try:\n",
        "            # Carregar o DataFrame\n",
        "            df = pd.read_csv(file_path)  # ou pd.read_excel(file_path) para arquivos Excel\n",
        "\n",
        "            # Renomear colunas se necessário\n",
        "            if 'label_included' in df.columns:\n",
        "                df.rename(columns={'label_included': 'label'}, inplace=True)\n",
        "            if 'labels' in df.columns:\n",
        "                df.rename(columns={'labels': 'label'}, inplace=True)\n",
        "\n",
        "            # Adicionar o caminho e nome do arquivo como colunas extras\n",
        "            df['file_path'] = file_path\n",
        "            df['file_name'] = os.path.basename(file_path)\n",
        "\n",
        "            # Adicionar o DataFrame à lista\n",
        "            dfs.append(df)\n",
        "            file_paths.append(file_path)\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f\"Erro ao processar o arquivo {file_path}: {str(e)}\")\n",
        "\n",
        "    # Concatenar todos os DataFrames em um único DataFrame\n",
        "    final_df = pd.concat(dfs, ignore_index=True)\n",
        "\n",
        "    # Imprimir informações sobre o DataFrame final\n",
        "    print(\"Informações sobre o DataFrame final:\")\n",
        "    print(final_df.info())\n",
        "\n",
        "    # Exemplo de visualização das primeiras linhas do DataFrame final\n",
        "    print(\"\\nPrimeiras linhas do DataFrame final:\")\n",
        "    print(final_df.head())\n",
        "\n",
        "    # Preencher valores nulos nas colunas 'title' e 'abstract'\n",
        "    for col in [\"title\", \"abstract\"]:\n",
        "        final_df.loc[final_df[col].isna(), col] = \"\"\n",
        "\n",
        "    # Criar a coluna 'text' concatenando 'title' e 'abstract'\n",
        "    final_df['text'] = final_df['title'] + \"; \" + final_df['abstract']\n",
        "\n",
        "    return final_df\n",
        "\n",
        "\n",
        "def model_train_validation(data_path, model_name, trys=1, config={}, just_notrain=False):\n",
        "    for i_try in range(trys):\n",
        "      ###### Init mdeol\n",
        "        ##########################################################################\n",
        "\n",
        "        model_ref_t = AutoModel.from_pretrained(model_name)\n",
        "        tokenizer_t_ref = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "        max_seq_length = 512\n",
        "\n",
        "        tokenize_batch_t = lambda train_text: tokenizer_t_ref.batch_encode_plus(\n",
        "            train_text,\n",
        "            add_special_tokens=True,\n",
        "            padding='max_length',\n",
        "            max_length=max_seq_length,\n",
        "            truncation=True,\n",
        "            return_tensors='pt',\n",
        "            return_attention_mask=True\n",
        "        )\n",
        "\n",
        "        data_name = os.path.basename(data_path).split(\".\")[0]\n",
        "\n",
        "        base = \"/content/gdrive/MyDrive/Projetos/TCC\"\n",
        "        base_save = f\"{base}/export/diagnostic/{model_name}/{data_name}\"\n",
        "        create_repo_if_not_exist(base_save+\"/\")\n",
        "        file_path = data_path\n",
        "\n",
        "\n",
        "        # model_ref_t, model_name\n",
        "\n",
        "\n",
        "      ###### Init data\n",
        "        ##########################################################################\n",
        "\n",
        "\n",
        "        df = preprocess_csv_file(file_path)\n",
        "\n",
        "        tk = lambda x: len(tokenizer_t_ref.encode(x))\n",
        "        tk_len = df.text.map(tk)\n",
        "\n",
        "        df_rest = df[tk_len > 20]\n",
        "\n",
        "        df_train = df_rest.groupby('label').sample(8)\n",
        "        df_test = df[~df.index.isin(df_train.index)]\n",
        "\n",
        "        exs_test = pd.concat([df_train, df_test])\n",
        "\n",
        "\n",
        "        sim_data = SIMDataset(df_train, tokenize_batch_t, config[\"n_pairs\"],\n",
        "                              # emb= emb, device=\"cuda\", pttn=\"* d a\", b=300\n",
        "                              )\n",
        "        train_dataloader = DataLoader(sim_data,shuffle=True,\n",
        "                                      batch_size=config[\"batch_size\"])\n",
        "        test_dataloader_eval = DataLoader(shuffle=False, batch_size=300,\n",
        "                    dataset=SIMDataset(\n",
        "                        exs_test, tokenize_batch_t, 0,\n",
        "                        # emb= emb, device=\"cuda\", pttn=\"* d a\"\n",
        "                              )\n",
        "                    )\n",
        "\n",
        "        train_dataloader_eval = DataLoader(shuffle=False, batch_size=len(df_train),\n",
        "                    dataset=SIMDataset(\n",
        "                        df_train, tokenize_batch_t, 0,\n",
        "                        # emb= emb, device=\"cuda\", pttn=\"* d a\"\n",
        "                              )\n",
        "                    )\n",
        "\n",
        "      ####\n",
        "        ##########################################################################\n",
        "        ########################    Nothing   ####################################\n",
        "        ##########################################################################\n",
        "\n",
        "        if just_notrain:\n",
        "\n",
        "          save = f\"{base_save}/No_train/{i_try}/\"\n",
        "          create_repo_if_not_exist(save)\n",
        "\n",
        "          smodel_mod_t = SentenceEmbedding(\n",
        "              copy.deepcopy(model_ref_t),\n",
        "              with_head = False)\n",
        "\n",
        "          smodel_mod_t.to(\"cuda\")\n",
        "          res = eval_pipeline(smodel_mod_t,\n",
        "                              train_dataloader=train_dataloader_eval,\n",
        "                              test_dataloader=test_dataloader_eval,\n",
        "                              device=\"cuda\",\n",
        "                              save= save)\n",
        "\n",
        "          export_json(res, f\"{save}Metrics_{i_try}.json\")\n",
        "\n",
        "          from pprint import pprint\n",
        "\n",
        "          pprint(res)\n",
        "\n",
        "\n",
        "          smodel_mod_t.to(\"cpu\")\n",
        "          torch.cuda.empty_cache()\n",
        "\n",
        "        else:\n",
        "\n",
        "        ##### modular\n",
        "          ##########################################################################\n",
        "          ########################    LoRA   #######################################\n",
        "          ##########################################################################\n",
        "\n",
        "          Mod = ModularUtils(Modular=ModMix)\n",
        "          Mod.freeze(module=model_ref_t)\n",
        "\n",
        "          model_base_t = copy.deepcopy(model_ref_t)\n",
        "          model_mod_t = Mod.copy_arch(model_base_t)\n",
        "\n",
        "          # model_mod_t.encoder.layer = model_mod_t.encoder.layer[-config[\"la\"]:]\n",
        "          # submod = model_mod_t.encoder.layer\n",
        "          submod = model_mod_t.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "          ###### model\n",
        "          smodel_mod_t = SentenceEmbedding(\n",
        "              model_mod_t,\n",
        "              with_head = True)\n",
        "\n",
        "\n",
        "          # Mod(smodel_mod_t)\n",
        "          Mod(smodel_mod_t)\n",
        "          smodel_mod_t.util.map(\n",
        "              name_module = [\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "\n",
        "          Mod(submod)\n",
        "          # activate in the last layers\n",
        "          submod.util.activate(\n",
        "              layer_names= [\"query\", 'value'],\n",
        "              key = 'set',\n",
        "              config = {\n",
        "                  'name':'test',\n",
        "                  **config\n",
        "                  })\n",
        "\n",
        "\n",
        "          smodel_mod_t.util.trainable_params()\n",
        "\n",
        "\n",
        "          l1 = loss_sim(smodel_mod_t,tau = config[\"tau\"], device=\"cuda\").to(\"cuda\")\n",
        "\n",
        "\n",
        "\n",
        "          # ##########################################################################\n",
        "          # ######## data   #################\n",
        "          # model_base_t2 = copy.deepcopy(model_ref_t)\n",
        "          # model_base_t2.encoder.layer = model_base_t2.encoder.layer[:config[\"rest\"]]\n",
        "\n",
        "          # model_base_t2.to(\"cuda\")\n",
        "          # model_base_t2.eval()\n",
        "          # emb = lambda **x: model_base_t2(**x)[0]\n",
        "\n",
        "\n",
        "\n",
        "          interation = int(len(sim_data)/config[\"batch_size\"])\n",
        "          total_iterations = config[\"n_epocs\"]*(interation)\n",
        "          ##########################################################################\n",
        "\n",
        "\n",
        "          def lr_lambda(epoch):\n",
        "              return 1 - (epoch / total_iterations)**2\n",
        "\n",
        "\n",
        "          # Create the linear learning rate scheduler\n",
        "          optimizer = Adam(smodel_mod_t.parameters(), lr=config[\"lr\"])\n",
        "          scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)\n",
        "\n",
        "          #### train\n",
        "          trainer = Trainer(losses=l1.to('cuda'),\n",
        "                            optimizer=optimizer,\n",
        "                            scheduler=scheduler,\n",
        "                            accumulation_steps = config[\"accumulation_steps\"],\n",
        "                            train_dataloader=train_dataloader,\n",
        "                            total_iterations=total_iterations,\n",
        "                            device = 'cuda')\n",
        "\n",
        "\n",
        "\n",
        "          print(f\"Quantidade de amostras\\n {len(sim_data)}\")\n",
        "          print(f\"Quantidade de exemplos\\n {len(df_train)}\")\n",
        "\n",
        "          w = submod[0].attention.self.value.Mix['Adapt'].loraA\n",
        "\n",
        "          print(submod)\n",
        "          print()\n",
        "          # print(w.weight)\n",
        "\n",
        "          trainer.train()\n",
        "\n",
        "          # print(w.weight)\n",
        "          ##########################################################################\n",
        "\n",
        "          # model_base_t2.to(\"cpu\")\n",
        "\n",
        "\n",
        "\n",
        "          save = f\"{base_save}/Lora/{i_try}/\"\n",
        "          create_repo_if_not_exist(save)\n",
        "\n",
        "\n",
        "          res = eval_pipeline(smodel_mod_t,\n",
        "                              train_dataloader=train_dataloader_eval,\n",
        "                              test_dataloader=test_dataloader_eval,\n",
        "                              device=\"cuda\",\n",
        "                              save= save)\n",
        "\n",
        "          export_json(res, f\"{save}Metrics_{i_try}.json\")\n",
        "\n",
        "          from pprint import pprint\n",
        "\n",
        "          pprint(res)\n",
        "\n",
        "\n",
        "          smodel_mod_t.to(\"cpu\")\n",
        "          torch.cuda.empty_cache()\n",
        "\n",
        "        ##### modular\n",
        "          ##########################################################################\n",
        "          ########################    QLoRA   ######################################\n",
        "          ##########################################################################\n",
        "          # model_base_t = copy.deepcopy(model_ref_t)\n",
        "\n",
        "\n",
        "          ##########################################################################\n",
        "          # ## Quant\n",
        "          smodel_mod_t.util.activate(\n",
        "                  layer_names= [\"query\", 'value']+[\"key\", 'dense'],\n",
        "                  key = 'set',\n",
        "                  config = {\n",
        "                      'name': 'Quant',\n",
        "                      'bits': 4,\n",
        "                      \"copyw\":True,\n",
        "                      \"base\":\"base_linear\"\n",
        "                      })\n",
        "\n",
        "          smodel_mod_t.util.activate(\n",
        "                  layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "                  key = 'quantize',\n",
        "                  Modular=Quantize_Linear)\n",
        "\n",
        "          ##########################################################################\n",
        "\n",
        "          ## lora\n",
        "          submod.util.activate(\n",
        "                  layer_names= [\"query\", 'value'],\n",
        "                  key = 'set', config= {'name': 'test',\n",
        "                                        'base': 'Quant'}\n",
        "          )\n",
        "\n",
        "\n",
        "          ##########################################################################\n",
        "\n",
        "          save = f\"{base_save}/Quant/{i_try}/\"\n",
        "          create_repo_if_not_exist(save)\n",
        "\n",
        "          smodel_mod_t.cuda()\n",
        "          smodel_mod_t.eval()\n",
        "          res = eval_pipeline(smodel_mod_t,\n",
        "                              train_dataloader=train_dataloader_eval,\n",
        "                              test_dataloader=test_dataloader_eval,\n",
        "                        device=\"cuda\",\n",
        "                              save=save)\n",
        "\n",
        "          from pprint import pprint\n",
        "          pprint(res)\n",
        "          export_json(res, f\"{save}Metrics_{i_try}.json\")\n",
        "\n",
        "          smodel_mod_t.to(\"cpu\")\n",
        "          torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "\n",
        "        ##### modular\n",
        "          ##########################################################################\n",
        "          ########################    SPARSE   #####################################\n",
        "          ##########################################################################\n",
        "\n",
        "          smodel_mod_t.util.activate(\n",
        "                  layer_names= [\"query\", 'value', \"key\"],\n",
        "                  key = 'set',\n",
        "                  config = {\n",
        "                      'name': 'Prune',\n",
        "                      's': 0.6,\n",
        "                      'acum': True,\n",
        "                      \"copyw\":True\n",
        "                      })\n",
        "\n",
        "\n",
        "          smodel_mod_t.util.activate(\n",
        "                  layer_names= ['dense'],\n",
        "                  key = 'set',\n",
        "                  config = {\n",
        "                      'name': 'Prune',\n",
        "                      's': 0.6,\n",
        "                      'acum': True,\n",
        "                      \"copyw\":True\n",
        "                      })\n",
        "\n",
        "          smodel_mod_t.to(\"cuda\")\n",
        "          get_batch_labels(\n",
        "              smodel_mod_t,\n",
        "              train_dataloader_eval)\n",
        "          smodel_mod_t.to(\"cpu\")\n",
        "\n",
        "\n",
        "          smodel_mod_t.util.activate(\n",
        "                  layer_names= ['dense']+[\"query\", 'value', \"key\"],\n",
        "                  key = 'prune',\n",
        "                  Modular=Prune\n",
        "                  )\n",
        "\n",
        "          smodel_mod_t.util.activate(\n",
        "                  layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "                  key = 'set',\n",
        "                  config = {\n",
        "                      'name': 'Prune',\n",
        "                      'fake_prune': True,\n",
        "                      })\n",
        "\n",
        "          ## lora\n",
        "          submod.util.activate(\n",
        "                  layer_names= [\"query\", 'value'],\n",
        "                  key = 'set', config= {'name': 'test',\n",
        "                                        'base': 'Prune'\n",
        "                                        }\n",
        "          )\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "          smodel_mod_t.cuda()\n",
        "          smodel_mod_t.eval()\n",
        "\n",
        "\n",
        "          save = f\"{base_save}/Prune/{i_try}/\"\n",
        "          create_repo_if_not_exist(save)\n",
        "\n",
        "\n",
        "          res = eval_pipeline(smodel_mod_t,\n",
        "                              train_dataloader=train_dataloader_eval,\n",
        "                              test_dataloader=test_dataloader_eval,\n",
        "                        device=\"cuda\",\n",
        "                              save=save)\n",
        "          smodel_mod_t.to(\"cpu\")\n",
        "          torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "          ########\n",
        "          m = copy.deepcopy(smodel_mod_t)\n",
        "\n",
        "          m.util.activate(\n",
        "                  layer_names= ['dense']+[\"query\", 'value', \"key\"],\n",
        "                  key = 'prune',\n",
        "                  Modular=Prune\n",
        "                  )\n",
        "          m.model.encoder.layer[0].attention.self.query.Mix.Prune.main.weight\n",
        "\n",
        "          sub_copy = m.model.encoder.layer[-3:]\n",
        "\n",
        "          ## lora\n",
        "          Mod(sub_copy)\n",
        "          m.util.activate(\n",
        "                  layer_names= [\"key\", 'dense'] + [\"query\", 'value'],\n",
        "                  key = 'del_layer',\n",
        "                  config={\n",
        "                      \"name\":[\"base_linear\", \"Quant\"],\n",
        "                      \"base\":\"Prune\"\n",
        "                  })\n",
        "\n",
        "          sub_copy.util.activate(\n",
        "                  layer_names= [\"query\", 'value'],\n",
        "                  key = 'set', config= {'name': 'test',\n",
        "                                        'base': 'Prune'}\n",
        "          )\n",
        "\n",
        "          sub_copy.util.activate(\n",
        "                  layer_names= [\"query\", 'value'],\n",
        "                  key = 'del_layer',\n",
        "                  config={\n",
        "                      \"name\":[\"Adapt\"],\n",
        "                      \"base\":\"Prune\"\n",
        "                  })\n",
        "          ##########\n",
        "\n",
        "\n",
        "          res[\"sparse\"] = 100 - count_non_zero_params(m)\n",
        "          export_json(res, f\"{save}Metrics_{i_try}.json\")\n",
        "\n",
        "          del m\n",
        "\n",
        "          pprint(res)\n",
        "\n",
        "\n",
        "        ##### modular\n",
        "          ##########################################################################\n",
        "          ##########################################################################\n",
        "          ##########################################################################\n",
        "\n",
        "          submod.util.activate(\n",
        "              layer_names= [\"query\", 'value'],\n",
        "              key = 'export',\n",
        "              Modular=Lora)\n",
        "\n",
        "          save = f\"{base_save}/Lora/{i_try}/\"\n",
        "          export_model(submod.util.export_mod(Modular=Lora),\n",
        "                      f\"{save}/lora_w.bin\")\n"
      ],
      "metadata": {
        "id": "AiVV4qkk6Xa_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os\n",
        "from pprint import pprint\n",
        "base_path = '/content/content/slr_data_tcc'\n",
        "# Paths and names of the domains datasets\n",
        "# path = 'SLR_data'\n",
        "paths1 = glob.glob(f\"{base_path}/SWIFT systematic review data/*.csv\", recursive=True)\n",
        "paths2 = glob.glob(f\"{base_path}/cohen/*.csv\", recursive=True)\n",
        "paths=[*paths1, *paths2]\n",
        "\n",
        "\n",
        "# names = [os.path.basename(p).split(\".\")[0] for p in paths]"
      ],
      "metadata": {
        "id": "NKlVsrBC1_A6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = {\n",
        "    \"total_batch\": 16,\n",
        "    \"accumulation_steps\": 1,\n",
        "    \"batch_size\":int(16/1),\n",
        "###\n",
        "    \"n_epocs\": 2,\n",
        "    \"n_pairs\" : 40,\n",
        "    \"lr\" : 7e-3,\n",
        "###\n",
        "    \"alpha\":4*(0.1),\n",
        "    \"r\":4,\n",
        "    \"tau\" : 0.2,\n",
        "###\n",
        "    \"la\" : 3,\n",
        "    \"rest\" : 12 - 3,\n",
        "    \"samples\":8\n",
        "           }\n",
        "\n",
        "seed = 43\n",
        "np.random.seed(seed)\n",
        "\n",
        "# Set seed for PyTorch\n",
        "torch.manual_seed(seed)\n",
        "\n",
        "# Set seed for CUDA if available\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed(seed)\n",
        "\n",
        "\n",
        "# base_dataset = \"cohen/Opiods.csv\"\n",
        "model_name = 'allenai/specter2_base'\n",
        "    # model_name = 'sentence-transformers/paraphrase-MiniLM-L3-v2'\n",
        "    # model_name = 'sentence-transformers/all-MiniLM-L12-v2'\n",
        "    # model_name = 'allenai/scibert_scivocab_uncased'\n",
        "\n",
        "# for p in paths[1:-1]:\n",
        "for p in paths:\n",
        "  model_train_validation(p, model_name, config=config, just_notrain = False)"
      ],
      "metadata": {
        "id": "ZZKOuXhbEoFG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### *metrics*"
      ],
      "metadata": {
        "id": "Q4GvrxpHWb-B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os\n",
        "from pprint import pprint\n",
        "base_path = '/content/gdrive/MyDrive/Projetos/TCC/export/diagnostic/allenai/specter2_base'\n",
        "# Paths and names of the domains datasets\n",
        "# path = 'SLR_data'\n",
        "\n",
        "paths_notrain = glob.glob(f\"{base_path}/**/No_train/0/*.json\", recursive=True)\n",
        "paths_lora = glob.glob(f\"{base_path}/**/Lora/0/*.json\", recursive=True)\n",
        "paths_quant = glob.glob(f\"{base_path}/**/Quant/0/*.json\", recursive=True)\n",
        "paths_prune = glob.glob(f\"{base_path}/**/Prune/0/*.json\", recursive=True)\n",
        "# pprint(paths)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# names = [os.path.basename(p).split(\".\")[0] for p in paths]"
      ],
      "metadata": {
        "id": "cLtP8XBWWfA8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "names_r = {\n",
        "    \"Opiods\": \"Opioids\",\n",
        "    \"ACEInhibitors\": \"ACE Inhibitors\",\n",
        "    \"AtypicalAntipsychotics\": \"Atypical Antipsychotics\",\n",
        "    \"BetaBlockers\": \"Beta Blockers\",\n",
        "    \"UrinaryIncontinence\": \"Urinary Incontinence\",\n",
        "    \"OralHypoglycemics\": \"Oral Hypoglycemics\",\n",
        "    \"ProtonPumpInhibitors\": \"Proton Pump Inhibitors\",\n",
        "    \"CalciumChannelBlockers\": \"Calcium Channel Blockers\",\n",
        "    \"SkeletalMuscleRelaxants\": \"Skeletal Muscle Relaxants\",\n",
        "    \"SkeletalMuscleRelaxants\": \"Skeletal Muscle Relaxants\"\n",
        "    }"
      ],
      "metadata": {
        "id": "VpFIdN1Vwvje"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "def json_data(paths):\n",
        "  metrics = []\n",
        "  for p in paths:\n",
        "    da = read_json(p)\n",
        "    for t in [\"95\", \"observed\"]:\n",
        "      metric = create_subset(da[t], [\"wss\", \"awss\", \"f1\", \"R\", \"prox.\"])\n",
        "      metric[\"method\"] = p.split(\"/\")[-3]\n",
        "      metric[\"dataset\"] = p.split(\"/\")[-4]\n",
        "      metric[\"type\"] = t\n",
        "      metrics.append(metric)\n",
        "\n",
        "  return metrics\n",
        "\n",
        "data_no_train = pd.DataFrame(json_data(paths_notrain))\n",
        "data_lora = pd.DataFrame(json_data(paths_lora))\n",
        "data_quant = pd.DataFrame(json_data(paths_quant))\n",
        "data_prune = pd.DataFrame(json_data(paths_prune))\n",
        "\n",
        "data_no_train[\"dataset\"] = data_no_train[\"dataset\"].replace(names_r)\n",
        "data_lora[\"dataset\"] = data_lora[\"dataset\"].replace(names_r)\n",
        "data_quant[\"dataset\"] = data_quant[\"dataset\"].replace(names_r)\n",
        "data_prune[\"dataset\"] = data_prune[\"dataset\"].replace(names_r)\n",
        "\n",
        "\n",
        "\n",
        "sub_no_train = data_no_train[data_no_train[\"type\"] == \"95\"]\n",
        "sub_lora = data_lora[data_lora[\"type\"] == \"95\"]\n",
        "sub_quant = data_quant[data_quant[\"type\"] == \"95\"]\n",
        "sub_prune = data_prune[data_prune[\"type\"] == \"95\"]\n",
        "\n",
        "\n",
        "\n",
        "plt.figure(figsize=(5, 8))\n",
        "\n",
        "plt.grid(True)\n",
        "plt.scatter( sub_lora['awss'],sub_lora['dataset'], color='blue',label=\"Lora\", alpha=0.9,marker=\"3\", s=350)\n",
        "plt.scatter( sub_quant['awss'],sub_quant['dataset'], color='green', label=\"Quant. 4bit\", alpha=0.9,  marker=\"3\", s=350)\n",
        "plt.scatter( sub_prune['awss'],sub_prune['dataset'], color='red', label=\"Prune\", alpha=0.9, marker=\"3\", s=350)\n",
        "plt.scatter( sub_no_train['awss'],sub_no_train['dataset'], color='black',label=\"Base\", alpha=1, marker='*', s=60)\n",
        "# plt.ylabel('Dataset', fontsize=14)\n",
        "plt.xlabel('AWSS', fontsize=14, weight=\"bold\")\n",
        "plt.xticks(fontsize=12)  # Set font size for x-axis ticks\n",
        "plt.yticks(fontsize=14,  weight=\"bold\")  # Set font size for y-axis ticks\n",
        "plt.legend(fontsize=11)\n",
        "# plt.title('Scatter Plot of AWSS by Dataset')\n",
        "plt.savefig('all.pdf')\n",
        "plt.show()\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# print(lora_data.round(3)[[\"awss\", \"dataset\", \"type\", \"R\"]].pivot_table(index='dataset', columns='type').style.to_latex())\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "5oysD_d2W4go"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## svd"
      ],
      "metadata": {
        "id": "vOvw9c6umAv1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "# Define the input size and output size for the linear layers\n",
        "input_size = 2\n",
        "output_size = 4\n",
        "\n",
        "# Define the first linear layer\n",
        "linear_layer1 = nn.Linear(input_size, output_size)\n",
        "\n",
        "# Define the second linear layer\n",
        "linear_layer2 = nn.Linear(input_size, output_size)\n",
        "\n",
        "delta= linear_layer1.weight - linear_layer2.weight\n",
        "\n",
        "# Concatenate the weights of the two linear layers\n",
        "\n",
        "w1 = linear_layer1.weight\n",
        "w2 = linear_layer2.weight\n",
        "\n",
        "wm = torch.cat((w1, w2), dim=1)\n",
        "\n",
        "linear_layer3 = nn.Linear(2 * input_size, output_size, bias=False)\n",
        "\n",
        "with torch.no_grad():\n",
        "    linear_layer3.weight.copy_(wm)\n",
        "\n",
        "\n",
        "linear_layer3.weight, w1, w2\n"
      ],
      "metadata": {
        "id": "HgKB41A9mJdX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a =0.5\n",
        "\n",
        "\n",
        "m1 = w1.norm(p=2, dim=0, keepdim=True)\n",
        "m2 = w2.norm(p=2, dim=0, keepdim=True)\n",
        "\n",
        "delta = a*w1/m1 + (1-a)*w2/m2\n",
        "delta\n",
        "# m = delta.norm(p=2, dim=0, keepdim=True)\n",
        "\n",
        "# delta_v = delta / m"
      ],
      "metadata": {
        "id": "84APMbGswFgQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "r=2\n",
        "U, S, V = torch.svd(delta)\n",
        "\n",
        "At = U[:, :r]\n",
        "Bt = torch.matmul(torch.diag(S[:r]), V[:r, :])\n",
        "\n",
        "w = torch.matmul(At, Bt)\n",
        "\n",
        "w*(a*m1+ (1-a)*m2)\n",
        "\n"
      ],
      "metadata": {
        "id": "03Aens0SpoKk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "-4xc5rWIpL4M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# modular\n",
        "Mod = ModularUtils(Modular=ModMix)\n",
        "\n",
        "# base model and the lora adapter model\n",
        "Mod.freeze(module=model_ref_t)\n",
        "\n",
        "\n",
        "model_base_t = copy.deepcopy(model_ref_t)\n",
        "model_base_t = SentenceEmbedding(\n",
        "    model_ref_t,\n",
        "    with_head = True\n",
        "    )\n",
        "\n",
        "model_mod_t = Mod.copy_arch(smodel_mod_t)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# save = f\"{base_save}/NoTrain/\"\n",
        "# create_repo_if_not_exist(save)\n",
        "\n",
        "\n",
        "# smodel_mod_t.eval().to(\"cuda\")\n",
        "# res = eval_pipeline(smodel_mod_t,\n",
        "#                     train_dataloader = DataLoader(\n",
        "#                       SIMDataset(df_train, tokenize_batch_t),\n",
        "#                       shuffle=False, batch_size=len(df_train)),\n",
        "#                     test_dataloader = DataLoader(\n",
        "#                       SIMDataset(exs_test, tokenize_batch_t),\n",
        "#                       shuffle=False, batch_size=300),\n",
        "#               device=\"cuda\", perplexity=30,\n",
        "#               save= save)\n",
        "\n",
        "\n",
        "# from pprint import pprint\n",
        "# pprint(res)\n",
        "\n",
        "# export_json(res, f\"{save}Metrics.json\")\n",
        "\n",
        "# smodel_mod_t.to(\"cpu\")\n",
        "\n",
        "# del smodel_mod_t, model_base_t\n",
        "# torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "OrTKn1xSmDU6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Main pipeline"
      ],
      "metadata": {
        "id": "isCeG4KH6hqn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "2.4 . Estimativa de recall\n",
        "Aqui descrevemos um método estatístico simples para estimar a recordação dada uma lista classificada de referências. Assumimos que um avaliador começa a triagem no topo da lista, fornecendo sequencialmente um rótulo para cada item: “Incluído” ou “Excluído”. Estamos interessados ​​em estimar quando o avaliador atingiu um determinado nível de recordação (digamos, 95%) para que ele/ela possa interromper a triagem e perceber os benefícios da classificação de documentos (ao mesmo tempo que mantém a confiança de que a maioria dos documentos relevantes foram descobertos). ).\n",
        "\n",
        "O método funciona examinando a extensão de intervalos consecutivos de documentos excluídos que ocorrem entre cada documento relevante durante a triagem. Os comprimentos desses vãos fornecem uma base para estimar a probabilidade local de relevância do documento. Se os documentos fossem selecionados em ordem aleatória, esses comprimentos de intervalo seguiriam uma distribuição geométrica com uma taxa (“probabilidade de sucesso”) refletindo a frequência subjacente de documentos relevantes (incluídos) na lista de referência. Da mesma forma, se considerarmos os intervalos entre um conjunto de mais de dois documentos consecutivos incluídos, então o comprimento total do intervalo, uma soma de variáveis ​​aleatórias geométricas independentes, seguiria a distribuição binomial negativa. Portanto, os comprimentos observados dos documentos excluídos podem ser usados ​​para estimar o parâmetro de taxa subjacente para os documentos incluídos. Esta taxa pode então, por sua vez, ser utilizada para estimar o número de documentos relevantes restantes na lista de documentos não rastreados. Na prática, à medida que o avaliador avança pela lista classificada de referências, as lacunas entre os documentos relevantes tenderão a aumentar em extensão, porque, por definição, as listas classificadas são carregadas antecipadamente com documentos relevantes. Por outras palavras, embora a distribuição binomial negativa presuma que os “sucessos” são distribuídos de forma independente e aleatória, neste cenário, a taxa de sucesso deve ser decrescente (em vez de fixa ou crescente). Como resultado, este método tende a resultar em uma estimativa conservadora de recordação. Ou seja, a recordação verdadeira obtida é muitas vezes ligeiramente superior ao seu valor estimado.\n",
        "\n",
        "Este algoritmo possui um parâmetro de “ajuste” que chamamos de “lookback”, denotado por δ . Este parâmetro determina o número de vãos considerados na estimativa da taxa de inclusão. Totalizamos a distância do vão, D , (em número de documentos) entre a posição de triagem atual e o δ -ésimo documento incluído anteriormente. Se a taxa de inclusão para documentos restantes for p , e supondo (hipoteticamente) que os documentos foram amostrados aleatoriamente para triagem, então\n",
        "\n",
        "Com essas informações, o recall estimado é calculado com base em\n",
        "do seguinte modo:\n",
        "onde\n",
        "é uma estimativa de p com base no D observado , TP é o número de documentos relevantes identificados até o momento pelo rastreador e U é o número de documentos restantes não rastreados.\n",
        "\n",
        "3.2.2 . Efeito do comprimento da lista\n",
        "Na Tabela A.4 , mostramos o custo resultante de diversas simulações utilizando a distribuição de pontuação BPA. Para cada simulação, a taxa geral de inclusão foi definida em p  = 0,15 e o lookback, δ  = 2. Usando esses parâmetros, simulamos listas classificadas de vários comprimentos. Os custos resultantes são calculados em média em 5 tentativas. Conforme mostrado na tabela, listas mais longas estão associadas a custos menores. Isto porque, para uma determinada taxa de inclusão local, existe um número fixo de documentos que devem ser analisados ​​para se obter uma boa estimativa dessa taxa. À medida que o comprimento da lista aumenta, este número fixo torna-se menor como uma fração da carga total de triagem.\n",
        "\n",
        "3.2.3 . Efeito da taxa de inclusão\n",
        "A Tabela A.5 mostra o efeito da taxa de inclusão global, p , no custo. Para cada simulação, o comprimento da lista foi fixado em 10.000 documentos e o lookback, δ  = 2. Os custos resultantes são calculados em média em 5 tentativas. Em geral, quanto menor for a taxa global de inclusão, maior será o custo. Para estimar uma pequena taxa de inclusão é necessário triar mais documentos, o que por sua vez aumenta o custo.\n",
        "\n",
        "3.2.4 . Efeito do lookback δ\n",
        "A Tabela A.6 mostra o efeito da variação sistemática do lookback, δ , nos dados simulados a partir do conjunto de dados de dor neuropática. Para esses experimentos, o comprimento da lista foi fixado em 30.000 e a taxa de inclusão, p , foi fixada em 0,17. Os custos resultantes foram calculados em média em 5 tentativas. A tabela mostra que, em geral, o custo é uma função crescente de δ . Observe que para δ  = 1 e δ  = 2, o custo foi realmente negativo, indicando que a recordação verdadeira obtida está na verdade abaixo dos 95% desejados. Conseqüentemente, escolher valores de δ muito pequenos pode levar a subestimar a recuperação.\n",
        "\n",
        "A Figura A.6 ilustra a relação entre a variabilidade da estimativa de recall (e o custo correspondente) e o parâmetro δ . No Painel (a) δ  = 1 e no Painel (b) δ  = 100. Observe a variabilidade drasticamente diminuída na estimativa de recall para o valor mais alto de δ . Em geral, aumentar δ diminui a variabilidade da estimativa de recordação. Portanto, há um trade-off entre a variabilidade do custo e o valor esperado do custo. Ignorando a variabilidade, a configuração ideal para δ é, idealmente, o menor valor tal que o custo esperado seja não negativo. Em geral, δ maior pode ser ideal quando p for maior. Por outro lado, quando p é muito baixo ou a lista é curta, pode não ser aconselhável usar um δ grande ."
      ],
      "metadata": {
        "id": "CIMBkun6MHkO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Negative-binomial distribuition**\n",
        "\n",
        "$X\\sim NBin(\\delta,p)$\n",
        "\n",
        "- $\\delta$ - Número de sucessos na amostra\n",
        "- $p$ - Probabilidade de sucesso\n",
        "- $x_i$ - Quantas amostras para ter o ultimo sucesso\n",
        "\n",
        "$$\\begin{aligned}\n",
        "p(x={x_i})&= {{\\delta+{x_i}-1}\\choose{{x_i}}}(1-p)^{x_i} p^\\delta \\\\\n",
        "l(p;x)&\\propto {x_i}\\log(1-p)+ r \\log (p) \\\\\n",
        "l(p;x_1,...,x_n) &\\propto \\sum{x_i}\\log(1-p)+ n\\delta \\log (p) \\\\\n",
        "\\end{aligned}$$\n",
        "\n",
        "Portanto\n",
        "\n",
        "$$\\begin{aligned}\n",
        "{\\partial \\over \\partial p}l(p;x_1,...,x_n) &=  {-\\sum{x_i} \\over (1-p)} + {n\\delta \\over p}=0\\\\\n",
        "\\implies \\hat{p} =  {\\delta \\over \\delta+ \\bar{x}}\n",
        "\\end{aligned}$$\n",
        "\n",
        "\n",
        "\n",
        "Recall estimado com $\\hat{p}\\times U = FN$\n",
        "\n",
        "$$R={TP\\over TP +\\hat{p}\\times U}$$\n",
        "- U Número de documentos restantes não rastreados\n",
        "- TP Número de documentos Relevantes\n",
        "\n",
        "$$AWSS = {TN\\%} -(1-FP\\%)$$\n"
      ],
      "metadata": {
        "id": "spKh7slOEVjT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## tests"
      ],
      "metadata": {
        "id": "QLC5fdBjThp3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_index_of_ith_one(i, random_list, maximum=None):\n",
        "    \"\"\"\n",
        "    Generate a random list of 0s and 1s and return the index of the i-th occurrence of '1'.\n",
        "\n",
        "    Parameters:\n",
        "    i : int\n",
        "        The occurrence of '1' to find.\n",
        "    random_list : list\n",
        "        The list of 0s and 1s.\n",
        "\n",
        "    Returns:\n",
        "    int\n",
        "        The index of the i-th occurrence of '1'.\n",
        "        Returns -1 if i exceeds the number of '1's in the list.\n",
        "    \"\"\"\n",
        "    ones_count = 0\n",
        "    zeros_count = 0\n",
        "    last_one_index = -1\n",
        "    if maximum is None:\n",
        "      maximum = len(random_list)\n",
        "\n",
        "\n",
        "    for index, value in enumerate(random_list):\n",
        "        if index > maximum:\n",
        "          break\n",
        "\n",
        "        if value == 1:\n",
        "            ones_count += 1\n",
        "            if ones_count == i:\n",
        "                last_one_index = index\n",
        "                break\n",
        "        else:\n",
        "            zeros_count += 1\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    if last_one_index != -1:\n",
        "        return last_one_index, zeros_count, ones_count\n",
        "    else:\n",
        "        return -1, zeros_count, ones_count\n",
        "\n",
        "\n",
        "d = i_th_occurrence = 10\n",
        "U=length = 10000\n",
        "\n",
        "# U*0.01\n",
        "random_list = np.random.binomial(1,0.01,length)\n",
        "\n",
        "\n",
        "last_one, D, ones_count = get_index_of_ith_one(i_th_occurrence, random_list, 100)\n",
        "\n",
        "\n",
        "print(f\"The index of the {i_th_occurrence}-th occurrence of '1' is: {last_one+1}\")\n",
        "from pprint import pprint\n",
        "\n",
        "\n",
        "e=0.0001\n",
        "p = (ones_count+e)/(ones_count+e+D)\n",
        "\n",
        "d=ones_count\n",
        "TP = 10\n",
        "FN = p*(U-d)\n",
        "r = (TP+d)/(TP+d+FN)\n",
        "\n",
        "print(f\"FN: {round(FN)}\")\n",
        "print(f\"phat: {round(p,2)}\")\n",
        "print(f\"D: {D}\")\n",
        "print(f\"Ones: {ones_count}\")\n",
        "print(f\"recall: {r}\")\n",
        "\n",
        "TP = 10\n",
        "FN = 1\n",
        "\n",
        "TP/(TP + FN)\n",
        "\n",
        "# print()\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "2_-UgwWPPbJh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.linear_model import LogisticRegression\n",
        "import numpy as np\n",
        "from sklearn.metrics import recall_score\n",
        "\n",
        "# Gerando dados de exemplo\n",
        "np.random.seed(42)\n",
        "X = np.random.rand(1000, 1) * 10  # 1000 amostras entre 0 e 10\n",
        "y = (X[:, 0] > 5).astype(int)  # Rotulando como 1 se X > 5, senão 0\n",
        "\n",
        "# Inicializando modelo de regressão logística\n",
        "modelo = LogisticRegression(max_iter=1000)\n",
        "\n",
        "# Número inicial de amostras no conjunto de treinamento\n",
        "n = 2\n",
        "\n",
        "indices_amostras_treino = np.random.choice(range(1000), n, replace=False)\n",
        "\n",
        "X_treino = X[indices_amostras_treino]\n",
        "y_treino = y[indices_amostras_treino]\n",
        "\n",
        "delta = 2\n",
        "\n",
        "# p_hat = delta / (delta+ mean(x) )\n",
        "d = 1 # um sucesso\n",
        "x=200 # no 200 rank\n",
        "\n",
        "\n",
        "TP = 100\n",
        "U= 1000\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "p = d/(d+x)\n",
        "\n",
        "\n",
        "R = TP / (TP + p*U)\n",
        "\n",
        "print(p)\n",
        "print(R)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "for tentativa in range(5):\n",
        "\n",
        "    # Treinando o modelo\n",
        "    modelo.fit(X_treino, y_treino)\n",
        "\n",
        "    # ramain data\n",
        "    id_test = np.isin(range(len(X)), indices_amostras_treino)\n",
        "    X_test, y_test = X[~id_test] , y[~id_test]\n",
        "\n",
        "    # Obtendo probabilidades para todas as amostras\n",
        "    probabilidades = modelo.predict_proba(X_test)[:, 1]\n",
        "\n",
        "    preds = modelo.predict(X_test)\n",
        "\n",
        "    # Ordenando amostras com base nas probabilidades\n",
        "    indices_ordenados = np.argsort(probabilidades)[::-1]\n",
        "    next_sample = indices_ordenados[:n]\n",
        "    indices_amostras_treino = np.concatenate([indices_amostras_treino, next_sample])\n",
        "\n",
        "\n",
        "\n",
        "    U = len(y_test) # remain to label\n",
        "    TP = np.sum(y_treino)\n",
        "\n",
        "    x= get_index_of_ith_one(delta, y_test[indices_ordenados])\n",
        "\n",
        "    p = d/(\n",
        "        +x)\n",
        "    R = TP /(TP + p*U)\n",
        "\n",
        "    print(f\"phat: {p}\")\n",
        "    print(f\"over R: {R}\")\n",
        "\n",
        "    X_treino, y_treino = X[indices_amostras_treino], y[indices_amostras_treino]\n",
        "\n",
        "\n",
        "    print('X_treino:',len(X_treino))\n",
        "    acuracia = modelo.score(X_test, y_test)\n",
        "    recall = recall_score(y_test, preds)\n",
        "    print(f\"Tentativa {tentativa + 1}, Acurácia: {acuracia:.2f}\")\n",
        "    print(f\"Recall: {recall:.2f}\")\n",
        "\n",
        "    print(f\"####\"*100)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Após o loop, você pode usar o modelo treinado conforme necessário.\n"
      ],
      "metadata": {
        "id": "GawCRa61LhaW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## main"
      ],
      "metadata": {
        "id": "GMo1aG6xTj7G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "directory_path = 'data'\n",
        "\n",
        "# List files in the directory\n",
        "with os.scandir(directory_path) as entries:\n",
        "    file_paths = [entry.path for entry in entries if entry.is_file()]\n",
        "\n",
        "t = LazyLoadDataset(file_paths)\n",
        "\n",
        "DataLoader(t, batch_size=1, shuffle=False)\n",
        "\n",
        "t[0]"
      ],
      "metadata": {
        "id": "uWzuxrPV1AJO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import DataLoader\n",
        "\n",
        "def get_index_of_ith_one(i, random_list, maximum=None):\n",
        "    \"\"\"\n",
        "    Generate a random list of 0s and 1s and return the index of the i-th occurrence of '1'.\n",
        "\n",
        "    Parameters:\n",
        "    i : int\n",
        "        The occurrence of '1' to find.\n",
        "    random_list : list\n",
        "        The list of 0s and 1s.\n",
        "\n",
        "    Returns:\n",
        "    int\n",
        "        The index of the i-th occurrence of '1'.\n",
        "        Returns -1 if i exceeds the number of '1's in the list.\n",
        "    \"\"\"\n",
        "    ones_count = 0\n",
        "    zeros_count = 0\n",
        "    last_one_index = -1\n",
        "    if maximum is None:\n",
        "      maximum = len(random_list)\n",
        "\n",
        "\n",
        "    for index, value in enumerate(random_list):\n",
        "        if index > maximum:\n",
        "          break\n",
        "\n",
        "        if value == 1:\n",
        "            ones_count += 1\n",
        "            if ones_count == i:\n",
        "                last_one_index = index\n",
        "                break\n",
        "        else:\n",
        "            zeros_count += 1\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    if last_one_index != -1:\n",
        "        return last_one_index, zeros_count, ones_count\n",
        "    else:\n",
        "        return -1, zeros_count, ones_count\n",
        "\n",
        "\n",
        "def apply_modular_operations(model_ref_t, la, lora_alpha, r):\n",
        "    \"\"\"Apply ModularUtils operations.\"\"\"\n",
        "    # Initialize ModularUtils\n",
        "    Mod = ModularUtils(Modular=ModMix)\n",
        "\n",
        "    # Freeze base model and the lora adapter model\n",
        "    Mod.freeze(module=model_ref_t)\n",
        "\n",
        "    # Create a copy of the base model\n",
        "    model_base_t = copy.deepcopy(model_ref_t)\n",
        "\n",
        "    # Copy architecture using ModularUtils\n",
        "    model_mod_t = Mod.copy_arch(model_base_t)\n",
        "\n",
        "    # Create SentenceEmbedding with head\n",
        "    smodel_mod_t = SentenceEmbedding(model_mod_t, with_head=True)\n",
        "\n",
        "    # Get weights of the head layer\n",
        "    head_weights = smodel_mod_t.head[1].weight\n",
        "\n",
        "    # Apply ModMix in all corresponding values\n",
        "    Mod(smodel_mod_t)\n",
        "    smodel_mod_t.util.map(name_module=[\"query\", 'value'] + [\"key\", 'dense'])\n",
        "\n",
        "    # Active Lora just in some layers\n",
        "    # Get the last layers and apply the lora\n",
        "    submod = smodel_mod_t.get_submodule('encoder.layer')[-la:]\n",
        "    Mod(submod)\n",
        "\n",
        "    # Activate Lora in the last layers\n",
        "    submod.util.activate(\n",
        "        layer_names=[\"query\", 'value'],\n",
        "        key='set',\n",
        "        config={\n",
        "            'alpha': lora_alpha,\n",
        "            'r': r,\n",
        "            'name': 'test'\n",
        "        }\n",
        "    )\n",
        "\n",
        "    return smodel_mod_t\n",
        "\n",
        "\n",
        "def prepare_training(X_train, tokenize_func, n_pairs, batch_size, n_epochs, smodel, tau, lr, accumulation_steps):\n",
        "    temp = SIMDataset(X_train, tokenize_func, n_pairs)\n",
        "    train_dataloader = DataLoader(temp,\n",
        "                                  shuffle=True,\n",
        "                                  batch_size=batch_size)\n",
        "\n",
        "    interation = int(len(temp)/batch_size)\n",
        "    total_iterations = n_epocs*(interation)\n",
        "\n",
        "\n",
        "\n",
        "    trainer = Trainer(losses=loss_sim(smodel, tau = tau).to('cuda'),\n",
        "                      optimizer=Adam(smodel.parameters(), lr=lr),\n",
        "                      accumulation_steps = accumulation_steps,\n",
        "                      train_dataloader=train_dataloader,\n",
        "                      total_iterations=total_iterations,\n",
        "                      device = 'cuda').train()\n",
        "\n",
        "\n",
        "\n",
        "import torch\n",
        "import numpy as np\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "\n",
        "def train_logistic_regression(X_train, Y_train, X_test, n_pairs, smodel):\n",
        "    emd_dataloader = DataLoader(SIMDataset(X_train, tokenize_batch_t), batch_size=batch_size)\n",
        "\n",
        "    all_emd = []\n",
        "    for x in emd_dataloader:\n",
        "        all_emd.append(smodel(x))\n",
        "    emd_x = torch.cat(all_emd)\n",
        "\n",
        "    head_model = LogisticRegression(max_iter=1000)\n",
        "    head_model.fit(emd_x, Y_train)\n",
        "\n",
        "    probabilities = head_model.predict_proba(X_test)[:, 1]\n",
        "    preds = head_model.predict(X_test)\n",
        "\n",
        "    indices_ordered = np.argsort(probabilities)[::-1]\n",
        "    idx_train = np.concatenate([idx_train, indices_ordered[:n_pairs]])\n",
        "\n",
        "    X_train_new, Y_train_new = X_train[idx_train], Y_train[idx_train]\n",
        "\n",
        "    return X_train_new, Y_train_new, preds, probabilities\n",
        "\n",
        "import copy\n",
        "\n",
        "\n",
        "\n",
        "    # return trainer\n",
        "\n",
        "# Usage:\n",
        "# trainer = prepare_training(X_train, tokenize_batch_t, n_pairs, batch_size, n_epochs, smodel, tau, lr, accumulation_steps)\n",
        "\n",
        "\n",
        "smodel_test = copy.deepcopy(smodel)\n",
        "smodel_test=smodel_test.to(torch.bfloat16)\n",
        "submod = smodel_test.encoder.layer[-config[\"la\"]:]\n",
        "\n",
        "_ = main_pipe_train(smodel_test, submod, *lod_train1 ,config)\n",
        "# w_1 = export_adapt_head(smodel_test, submod)\n",
        "\n",
        "##### model\n",
        "\n",
        "n_pairs = 40\n",
        "n_epocs=2\n",
        "batch_size=16\n",
        "accumulation_steps=1\n",
        "tau=0.2\n",
        "lr=0.\n",
        "########################\n",
        "\n",
        "### gen seed\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "base_dataset5 = \"cohen/Opiods.csv\"\n",
        "file_path = f\"data/content/slr_data_tcc/{base_dataset5}\"\n",
        "df = preprocess_csv_file(file_path)\n",
        "\n",
        "\n",
        "\n",
        "X = df\n",
        "Y = df['label']\n",
        "idx_all = X.index\n",
        "\n",
        "seed = X.groupby('label').sample(4)\n",
        "\n",
        "df_train = seed\n",
        "idx_train = df_train.index\n",
        "df_test = df[~df.index.isin(idx_train)]\n",
        "\n",
        "X_train, Y_train = (df_train, df_train['label'])\n",
        "\n",
        "\n",
        "for tentativa in range(5):\n",
        "\n",
        "############ start here\n",
        "X_test = X[~idx_all.isin(idx_train)]\n",
        "Y_test = Y[~idx_all.isin(idx_train)]\n",
        "\n",
        "\n",
        "idx_test = df_test.index\n",
        "\n",
        "\n",
        "#####################\n",
        "# first train sbert\n",
        "\n",
        "lod1 = (train_dataloader_eval1, test_dataloader_eval1)\n",
        "lod_train1= (df_train1, train_dataloader, *lod1)\n",
        "\n",
        "\n",
        "\n",
        "temp = SIMDataset(X_train, tokenize_batch_t, n_pairs)\n",
        "train_dataloader = DataLoader(temp,\n",
        "                              shuffle=True,\n",
        "                              batch_size=batch_size)\n",
        "\n",
        "interation = int(len(temp)/batch_size)\n",
        "total_iterations = n_epocs*(interation)\n",
        "\n",
        "\n",
        "trainer = Trainer(losses=loss_sim(model, tau = tau).to('cuda'),\n",
        "                  optimizer=Adam(model.parameters(), lr=lr),\n",
        "                  accumulation_steps = accumulation_steps,\n",
        "                  train_dataloader=train_dataloader,\n",
        "                  total_iterations=total_iterations,\n",
        "                  device = 'cuda').train()\n",
        "\n",
        "#####################\n",
        "\n",
        "# second train sbert\n",
        "\n",
        "emd_dataloader = DataLoader(SIMDataset(X_train, tokenize_batch_t),\n",
        "                              batch_size)\n",
        "\n",
        "\n",
        "\n",
        "all = []\n",
        "for x in emd_dataloader:\n",
        "  all.append(model(x))\n",
        "\n",
        "emd_x =  torch.concat(all)\n",
        "\n",
        "head_model = LogisticRegression(max_iter=1000)\n",
        "head_model.fit(emd_x, Y_train)\n",
        "\n",
        "\n",
        "probabilidades = head_model.predict_proba(X_test)[:, 1]\n",
        "preds = modelo.predict(X_test)\n",
        "\n",
        "indices_ordenados = np.argsort(probabilidades)[::-1]\n",
        "idx_train = np.concatenate([idx_train, indices_ordenados[:n]])\n",
        "\n",
        "\n",
        "X_train, y_train = X[idx_train], Y[idx_train]\n",
        "\n",
        "#####################\n",
        "# metrics\n",
        "\n",
        "U = len(Y_test) # remain to label\n",
        "TP = np.sum(Y_train)\n",
        "\n",
        "x= get_index_of_ith_one(delta, Y_test[indices_ordenados])\n",
        "\n",
        "p = d/(d+x)\n",
        "R = TP /(TP + p*U)\n",
        "\n",
        "\n",
        "metrics(probabilidades, Y_test, recall_target=0.95)\n",
        "\n",
        "acuracia = modelo.score(X_test, Y_test)\n",
        "recall = recall_score(Y_test, preds)\n",
        "\n",
        "print(f\"Tentativa {tentativa + 1}, Acurácia: {acuracia:.2f}\")\n",
        "print(f\"Recall: {recall:.2f}\")\n"
      ],
      "metadata": {
        "id": "4k7LhfoFTdjm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Teste"
      ],
      "metadata": {
        "id": "dPUzUls7TWcu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install einops einop"
      ],
      "metadata": {
        "id": "rAaOROkPTY1J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import einops\n",
        "import torch\n",
        "import numpy as np\n",
        "from einop import einop\n",
        "\n",
        "\n",
        "y = torch.randn(size=(2,2,2,3))\n",
        "\n",
        "abs(y)"
      ],
      "metadata": {
        "id": "_a0SL00BTh1d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "einop(abs(y),\"b (i i0) j k->(b i) i0 (j k)\", i0=2)"
      ],
      "metadata": {
        "id": "hMaRMNPPld9N"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "einop(einop(abs(y),\"b i j k->(b i) (j k)\"),\n",
        "      \"(i i0) j -> i i0 j\", i0=2)"
      ],
      "metadata": {
        "id": "9ojFnJ6uRowd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "abs(y) / einop(abs(y),\"b m n d->b m n 1\", 'sum')"
      ],
      "metadata": {
        "id": "XFOni6qNlYvZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "einop(y,\"b m n d-> b (m n) d\")\n",
        "\n",
        "# abs(y) / einop(abs(y),\"b m n d->1 m n d\", 'sum')"
      ],
      "metadata": {
        "id": "v1K9e8K3Y78m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "abs(y) / einop(abs(y),\"b m n d->1 m n d\", 'sum')"
      ],
      "metadata": {
        "id": "ItIHI_riky1-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "einop(y,\"b m n d->b m d\", 'min')\n"
      ],
      "metadata": {
        "id": "yePdUVicgyfx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q einops"
      ],
      "metadata": {
        "id": "20gvXUmrRyVi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from einops import pack, unpack\n",
        "import numpy as np\n",
        "# h, w = 2, 2\n",
        "# image_rgb is 3-dimensional (h, w, 3) and depth is 2-dimensional (h, w)\n",
        "a = torch.Tensor(np.random.random([32,2]))\n",
        "b = torch.Tensor(np.random.random([32,2]))\n",
        "# but we can stack them\n"
      ],
      "metadata": {
        "id": "N56CsDJOp6vg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "c, ps = pack([a, b], '* b')\n",
        "\n",
        "print(c.shape)\n",
        "# c"
      ],
      "metadata": {
        "id": "stVnqNfuumbE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "d = unpack(c, [[1]]*7,\"* a b\")\n",
        "\n",
        "print(d[0].shape)\n",
        "d"
      ],
      "metadata": {
        "id": "j1FFcl_AqcXE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.einsum('i j k, i j k-> j k', d[0:2]).shape"
      ],
      "metadata": {
        "id": "mtfJMLMVxWPV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from einops import rearrange, reduce, repeat\n",
        "\n",
        "x = torch.randn((2,4,4))\n",
        "x"
      ],
      "metadata": {
        "id": "of3mj319yxVO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "reduce(x, 'b (h h2) (w w2) -> b h w', 'max', h2=2, w2=1).shape\n",
        "reduce(x, 'b (h h2) (w w2) -> b h w', 'max', h2=2, w2=1)"
      ],
      "metadata": {
        "id": "Y1SKlM9LcUhi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from einops.layers.torch import EinMix as Mix\n",
        "\n",
        "# other stuff we use\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "x = torch.randn((2,4,4))"
      ],
      "metadata": {
        "id": "R19gTHjFAIZl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "print(\"## x:\\n\", x)\n",
        "print(\"#####\"*3)\n",
        "\n",
        "a = Mix('b i j->b i i0', weight_shape='j i0',\n",
        "        # j =4,\n",
        "        # i =1,\n",
        "        i0 =2,\n",
        "        j =4,\n",
        "        # j0=2,\n",
        "        # k= 2,\n",
        "        )\n",
        "\n",
        "w = list(a.parameters())\n",
        "\n",
        "w[0].requires_grad=False\n",
        "\n",
        "w[0][:] = torch.Tensor([1])\n",
        "\n",
        "print(f\"### \\n(x) {x.shape}\\n(w) {w[0].shape} \\n(y) {a(x).shape}\")\n",
        "print(\"#####\"*3)\n",
        "print(\"## w:\\n\", w)\n",
        "print(\"#####\"*3)\n",
        "print(\"## aw(x):\\n\", a(x))\n",
        "x\n",
        "a"
      ],
      "metadata": {
        "id": "2xhWwsqOCkgv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# w[0].requires_grad=False\n",
        "\n",
        "\n",
        "# w[0][:,[0,3]] = torch.Tensor([0])\n",
        "\n",
        " (x[0][:1,:] * w[0]).sum()"
      ],
      "metadata": {
        "id": "pEmKX4TiUZ27"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.sum(x[0][:2,0:2]*w[0][0])\n",
        "torch.sum(x[0][2:,0:2]*w[0][0])\n",
        "# torch.sum(x[0][0,:]*w[0][0]) + torch.sum(x[0][1,:]*w[0][1])\n",
        "# x[0]*w[0][0,:,:]"
      ],
      "metadata": {
        "id": "Ul_yR7eOKVCf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "okLPyQTvb1td"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "print(\"## x:\\n\", x)\n",
        "print(\"#####\"*3)\n",
        "\n",
        "a = Mix('b i (j j0)-> b i', weight_shape='i j0',\n",
        "        i =2,\n",
        "        # j=2,\n",
        "        # i0=4,\n",
        "        j0=2,\n",
        "        )\n",
        "\n",
        "\n",
        "w = list(a.parameters())\n",
        "print(f\"### \\n(x) {x.shape}\\n(w) {w[0].shape} \\n(y) {a(x).shape}\")\n",
        "print(\"#####\"*3)\n",
        "print(\"## w:\\n\", w)\n",
        "print(\"#####\"*3)\n",
        "print(\"## aw(x):\\n\", a(x))\n"
      ],
      "metadata": {
        "id": "dj6o--xKCkfi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "# x = torch.randn((2,4,4))\n",
        "\n",
        "x = torch.Tensor([\n",
        "    [1,0],\n",
        "    [0,1],\n",
        "    ])\n",
        "\n",
        "a = torch.eye(2) * torch.Tensor([1,1])\n",
        "\n",
        "\n",
        "x.shape, a.shape\n",
        "\n",
        "torch.matmul(x, a)\n",
        "\n"
      ],
      "metadata": {
        "id": "Vfw_fFDcvpvU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x = torch.randn((2,2))\n",
        "\n",
        "\n",
        "a,b = torch.linalg.eig(x)\n",
        "a,b"
      ],
      "metadata": {
        "id": "Vht7PVgrwuWk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x"
      ],
      "metadata": {
        "id": "kntIZe9qxUrq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "b @ a"
      ],
      "metadata": {
        "id": "XBxn0q-JxPGQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}