BART Model Set Up.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -U  torch transformers apex torchsummary jupyter ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\stbyu\\anaconda3\\envs\\eecs595python10\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "import math\n",
    "import random\n",
    "import torch\n",
    "import warnings\n",
    "from torch import nn\n",
    "from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss\n",
    "from transformers import BartForConditionalGeneration, BartTokenizer, BartConfig\n",
    "from transformers.file_utils import (\n",
    "    add_code_sample_docstrings,\n",
    "    add_end_docstrings,\n",
    "    add_start_docstrings,\n",
    "    add_start_docstrings_to_model_forward,\n",
    "    replace_return_docstrings,\n",
    ")\n",
    "from transformers.modeling_outputs import (\n",
    "    BaseModelOutput,\n",
    "    BaseModelOutputWithPastAndCrossAttentions,\n",
    "    Seq2SeqLMOutput,\n",
    "    Seq2SeqModelOutput,\n",
    "    Seq2SeqQuestionAnsweringModelOutput,\n",
    "    Seq2SeqSequenceClassifierOutput,\n",
    ")\n",
    "from transformers.modeling_utils import PreTrainedModel\n",
    "from transformers.models.bart.modeling_bart import (\n",
    "    BartLearnedPositionalEmbedding,\n",
    "    BartDecoderLayer,\n",
    "    BartPreTrainedModel,\n",
    ")\n",
    "from transformers.utils import logging\n",
    "from typing import List, Optional, Tuple, Union\n",
    "from transformers.modeling_attn_mask_utils import (\n",
    "    _prepare_4d_attention_mask,\n",
    "    _prepare_4d_causal_attention_mask,\n",
    ")\n",
    "from transformers.activations import ACT2FN\n",
    "from tqdm import tqdm\n",
    "\n",
    "## Create a DataLoader\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "from transformers import AdamW"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BART Model Set Up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VectorizedFeatureSpecificMultiHeadAttention(torch.nn.Module):\n",
    "    def __init__(self, num_heads, d_model, num_features):\n",
    "        \"\"\"\n",
    "        Initializes the VectorizedFeatureSpecificMultiHeadAttention module.\n",
    "\n",
    "        Parameters:\n",
    "        num_heads (int): Number of attention heads.\n",
    "        d_model (int): Dimensionality of the input feature space.\n",
    "        num_features (int): Number of distinct features / experts.\n",
    "        \"\"\"\n",
    "        super(VectorizedFeatureSpecificMultiHeadAttention, self).__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.d_model = d_model\n",
    "        self.num_features = num_features\n",
    "\n",
    "        assert d_model % self.num_heads == 0, \"d_model must be divisible by num_heads\"\n",
    "\n",
    "        self.depth = d_model // self.num_heads\n",
    "\n",
    "        self.wq = torch.nn.Linear(d_model, d_model)\n",
    "        self.wk = torch.nn.Linear(\n",
    "            d_model, d_model\n",
    "        )  # Same as original since key is already feature-specific\n",
    "\n",
    "    def split_heads(self, x, batch_size, seq_len):\n",
    "        \"\"\"\n",
    "        Splits the last dimension of x into (num_heads, depth) and reshapes.\n",
    "\n",
    "        Parameters:\n",
    "        x (torch.Tensor): Input tensor.\n",
    "        seq_len (int): Sequence length for the reshaping process.\n",
    "\n",
    "        Returns:\n",
    "        torch.Tensor: Reshaped tensor.\n",
    "        \"\"\"\n",
    "        new_shape = x.size()[:-1] + (self.num_heads, self.depth)\n",
    "        x = x.reshape(*new_shape).permute(\n",
    "            0, 1, 3, 2, 4\n",
    "        )  # (num_features, batch_size, num_heads, seq_len, depth)\n",
    "        return x\n",
    "\n",
    "    def forward(self, query, key):\n",
    "        \"\"\"\n",
    "        Forward pass for the VectorizedFeatureSpecificMultiHeadAttention.\n",
    "\n",
    "        Parameters:\n",
    "        query (torch.Tensor): Query tensor of shape (batch_size, seq_len_q, d_model).\n",
    "        key (torch.Tensor): Key tensor of shape (num_features, batch_size, seq_len_k, d_model).\n",
    "\n",
    "        Returns:\n",
    "        torch.Tensor: Concatenated attention weights across all heads and features.\n",
    "        \"\"\"\n",
    "        batch_size = query.size(0)\n",
    "        seq_len_q = query.size(1)\n",
    "        seq_len_k = key.size(2)\n",
    "\n",
    "        # Prepare query and key\n",
    "        query = self.split_heads(\n",
    "            self.wq(query), batch_size, seq_len_q\n",
    "        )  # (batch_size, num_heads, seq_len_q, depth)\n",
    "        key = self.split_heads(\n",
    "            self.wk(key), batch_size, seq_len_k\n",
    "        )  # (num_features, batch_size, num_heads, seq_len_k, depth)\n",
    "\n",
    "        # Compute attention\n",
    "        attention_weights = self.compute_attention(query, key, seq_len_k)\n",
    "\n",
    "        return attention_weights\n",
    "\n",
    "    def compute_attention(self, query, key, seq_len_k):\n",
    "        \"\"\"\n",
    "        Computes the scaled dot-product attention.\n",
    "\n",
    "        Parameters:\n",
    "        query (torch.Tensor): Query tensor.\n",
    "        key (torch.Tensor): Key tensor.\n",
    "        seq_len_k (int): Key sequence length.\n",
    "\n",
    "        Returns:\n",
    "        torch.Tensor: Attention weights.\n",
    "        \"\"\"\n",
    "        # Expand query to match key's features\n",
    "        # query = query.unsqueeze(0).expand(\n",
    "        #     self.num_features, -1, -1, -1, -1\n",
    "        # )  # (num_features, batch_size, num_heads, seq_len_q, depth)\n",
    "\n",
    "        # Perform batch matrix multiplication\n",
    "        matmul_qk = torch.matmul(\n",
    "            query, key.transpose(-2, -1)\n",
    "        )  # (num_features, batch_size, num_heads, seq_len_q, seq_len_k)\n",
    "        \n",
    "        # Scale\n",
    "        scale_factor = 1 / torch.sqrt(torch.tensor(self.depth, dtype=torch.float32))\n",
    "        attention_weights = matmul_qk * scale_factor\n",
    "\n",
    "        # # Apply softmax along the dimension of key's sequence length\n",
    "        # attention_weights = F.softmax(scaled_attention_logits, dim=-1)\n",
    "\n",
    "        # Reshape to combine features and batch\n",
    "        # _, _, num_heads, _, _ = scaled_attention_logits.shape\n",
    "        # attention_weights = scaled_attention_logits.permute(\n",
    "        #     1, 2, 0, 3, 4\n",
    "        # ).contiguous()  # (batch_size, num_heads, num_features, seq_len_q, seq_len_k)\n",
    "\n",
    "        return attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def scaled_dot_product_attention(query, key, d_k):\n",
    "#     \"\"\"\n",
    "#     Compute the scaled dot product attention weights.\n",
    "\n",
    "#     Args:\n",
    "#     query (Tensor): Query tensor of shape (num_replicas, batch_size,seq_len_query,  d_k) or (seq_len_query, batch_size, d_k) if not batched.\n",
    "\n",
    "#     key (Tensor): Key tensor of shape (num_replicas,  batch_size, seq_len_key, d_k) or (batch_size, seq_len_key, d_k) if not batched.\n",
    "#     d_k (int): Dimension of the key and query tensors.\n",
    "\n",
    "#     Returns:\n",
    "#     Tensor: Attention weights of shape (num_replicas, seq_len_query, batch_size, seq_len_key)\n",
    "#             or (seq_len_query, batch_size, seq_len_key) if not batched.\n",
    "#     \"\"\"\n",
    "\n",
    "#     # Compute the dot product between query and key\n",
    "#     scores = torch.matmul(\n",
    "#         query, key.transpose(-2, -1)\n",
    "#     )  # Shape: (num_replicas, batch_size, seq_len_query, seq_len_key)\n",
    "\n",
    "#     # Scale the scores\n",
    "#     attention_weights = scores / torch.sqrt(torch.tensor(d_k, dtype=scores.dtype))\n",
    "#     # print(\"query shape: \", query.shape)\n",
    "#     # print(\"key shape: \", key.shape)\n",
    "#     # print(\"attention_weights shape: \", attention_weights.shape)\n",
    "\n",
    "#     return attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomBartDecoder(BartPreTrainedModel):\n",
    "    \"\"\"\n",
    "    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`BartDecoderLayer`]\n",
    "\n",
    "    Args:\n",
    "        config: BartConfig\n",
    "        embed_tokens (nn.Embedding): output embedding\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):\n",
    "        super().__init__(config)\n",
    "        self.dropout = config.dropout\n",
    "        self.layerdrop = config.decoder_layerdrop\n",
    "        self.padding_idx = config.pad_token_id\n",
    "        self.max_target_positions = config.max_position_embeddings\n",
    "        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0\n",
    "\n",
    "        self.embed_tokens = nn.Embedding(\n",
    "            config.vocab_size, config.d_model, self.padding_idx\n",
    "        )\n",
    "\n",
    "        if embed_tokens is not None:\n",
    "            self.embed_tokens.weight = embed_tokens.weight\n",
    "\n",
    "        self.embed_positions = BartLearnedPositionalEmbedding(\n",
    "            config.max_position_embeddings,\n",
    "            config.d_model,\n",
    "        )\n",
    "        self.layers = nn.ModuleList(\n",
    "            [BartDecoderLayer(config) for _ in range(config.decoder_layers)]\n",
    "        )\n",
    "        self.layernorm_embedding = nn.LayerNorm(config.d_model)\n",
    "\n",
    "        self.gradient_checkpointing = False\n",
    "        # Initialize weights and apply final processing\n",
    "        self.post_init()\n",
    "\n",
    "    def get_input_embeddings(self):\n",
    "        return self.embed_tokens\n",
    "\n",
    "    def set_input_embeddings(self, value):\n",
    "        self.embed_tokens = value\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        input_ids: torch.LongTensor = None,\n",
    "        attention_mask: Optional[torch.Tensor] = None,\n",
    "        encoder_hidden_states: Optional[torch.FloatTensor] = None,\n",
    "        encoder_attention_mask: Optional[torch.LongTensor] = None,\n",
    "        head_mask: Optional[torch.Tensor] = None,\n",
    "        cross_attn_head_mask: Optional[torch.Tensor] = None,\n",
    "        past_key_values: Optional[List[torch.FloatTensor]] = None,\n",
    "        inputs_embeds: Optional[torch.FloatTensor] = None,\n",
    "        use_cache: Optional[bool] = None,\n",
    "        output_attentions: Optional[bool] = None,\n",
    "        output_hidden_states: Optional[bool] = None,\n",
    "        return_dict: Optional[bool] = None,\n",
    "    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:\n",
    "        r\"\"\"\n",
    "        Args:\n",
    "            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):\n",
    "                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you\n",
    "                provide it.\n",
    "\n",
    "                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n",
    "                [`PreTrainedTokenizer.__call__`] for details.\n",
    "\n",
    "                [What are input IDs?](../glossary#input-ids)\n",
    "            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):\n",
    "                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:\n",
    "\n",
    "                - 1 for tokens that are **not masked**,\n",
    "                - 0 for tokens that are **masked**.\n",
    "\n",
    "                [What are attention masks?](../glossary#attention-mask)\n",
    "            encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):\n",
    "                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention\n",
    "                of the decoder.\n",
    "            encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):\n",
    "                Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values\n",
    "                selected in `[0, 1]`:\n",
    "\n",
    "                - 1 for tokens that are **not masked**,\n",
    "                - 0 for tokens that are **masked**.\n",
    "\n",
    "                [What are attention masks?](../glossary#attention-mask)\n",
    "            head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):\n",
    "                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:\n",
    "\n",
    "                - 1 indicates the head is **not masked**,\n",
    "                - 0 indicates the head is **masked**.\n",
    "\n",
    "            cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):\n",
    "                Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing\n",
    "                cross-attention on hidden heads. Mask values selected in `[0, 1]`:\n",
    "\n",
    "                - 1 indicates the head is **not masked**,\n",
    "                - 0 indicates the head is **masked**.\n",
    "\n",
    "            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):\n",
    "                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of\n",
    "                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of\n",
    "                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.\n",
    "\n",
    "                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the\n",
    "                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.\n",
    "\n",
    "                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those\n",
    "                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of\n",
    "                all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of\n",
    "                shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing\n",
    "                `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more\n",
    "                control over how to convert `input_ids` indices into associated vectors than the model's internal\n",
    "                embedding lookup matrix.\n",
    "            output_attentions (`bool`, *optional*):\n",
    "                Whether or not to return the attentions tensors of all attention layers. See `attentions` under\n",
    "                returned tensors for more detail.\n",
    "            output_hidden_states (`bool`, *optional*):\n",
    "                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors\n",
    "                for more detail.\n",
    "            return_dict (`bool`, *optional*):\n",
    "                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n",
    "        \"\"\"\n",
    "        output_attentions = (\n",
    "            output_attentions\n",
    "            if output_attentions is not None\n",
    "            else self.config.output_attentions\n",
    "        )\n",
    "        output_hidden_states = (\n",
    "            output_hidden_states\n",
    "            if output_hidden_states is not None\n",
    "            else self.config.output_hidden_states\n",
    "        )\n",
    "        use_cache = use_cache if use_cache is not None else self.config.use_cache\n",
    "        return_dict = (\n",
    "            return_dict if return_dict is not None else self.config.use_return_dict\n",
    "        )\n",
    "\n",
    "        # retrieve input_ids and inputs_embeds\n",
    "        if input_ids is not None and inputs_embeds is not None:\n",
    "            raise ValueError(\n",
    "                \"You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time\"\n",
    "            )\n",
    "        elif input_ids is not None:\n",
    "            input = input_ids\n",
    "            input_shape = input.shape\n",
    "            input_ids = input_ids.view(-1, input_shape[-1])\n",
    "        elif inputs_embeds is not None:\n",
    "            input_shape = inputs_embeds.size()[:-1]\n",
    "            input = inputs_embeds[:, :, -1]\n",
    "        else:\n",
    "            raise ValueError(\n",
    "                \"You have to specify either decoder_input_ids or decoder_inputs_embeds\"\n",
    "            )\n",
    "\n",
    "        # past_key_values_length\n",
    "        past_key_values_length = (\n",
    "            past_key_values[0][0].shape[2] if past_key_values is not None else 0\n",
    "        )\n",
    "\n",
    "        if inputs_embeds is None:\n",
    "            inputs_embeds = self.embed_tokens(input) * self.embed_scale\n",
    "\n",
    "        if getattr(self.config, \"_flash_attn_2_enabled\", False):\n",
    "            # 2d mask is passed through the layers\n",
    "            attention_mask = (\n",
    "                attention_mask\n",
    "                if (attention_mask is not None and 0 in attention_mask)\n",
    "                else None\n",
    "            )\n",
    "        else:\n",
    "            # 4d mask is passed through the layers\n",
    "            attention_mask = _prepare_4d_causal_attention_mask(\n",
    "                attention_mask, input_shape, inputs_embeds, past_key_values_length\n",
    "            )\n",
    "\n",
    "        # expand encoder attention mask\n",
    "        if encoder_hidden_states is not None and encoder_attention_mask is not None:\n",
    "            if getattr(self.config, \"_flash_attn_2_enabled\", False):\n",
    "                encoder_attention_mask = (\n",
    "                    encoder_attention_mask if 0 in encoder_attention_mask else None\n",
    "                )\n",
    "            else:\n",
    "                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n",
    "                encoder_attention_mask = _prepare_4d_attention_mask(\n",
    "                    encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]\n",
    "                )\n",
    "\n",
    "        # embed positions\n",
    "        positions = self.embed_positions(input, past_key_values_length)\n",
    "        positions = positions.to(inputs_embeds.device)\n",
    "\n",
    "        hidden_states = inputs_embeds + positions\n",
    "        hidden_states = self.layernorm_embedding(hidden_states)\n",
    "\n",
    "        hidden_states = nn.functional.dropout(\n",
    "            hidden_states, p=self.dropout, training=self.training\n",
    "        )\n",
    "\n",
    "        if self.gradient_checkpointing and self.training:\n",
    "            if use_cache:\n",
    "                logger.warning_once(\n",
    "                    \"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\"\n",
    "                )\n",
    "                use_cache = False\n",
    "\n",
    "        # decoder layers\n",
    "        all_hidden_states = () if output_hidden_states else None\n",
    "        all_self_attns = () if output_attentions else None\n",
    "        all_cross_attentions = (\n",
    "            () if (output_attentions and encoder_hidden_states is not None) else None\n",
    "        )\n",
    "        next_decoder_cache = () if use_cache else None\n",
    "\n",
    "        # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired\n",
    "        for attn_mask, mask_name in zip(\n",
    "            [head_mask, cross_attn_head_mask], [\"head_mask\", \"cross_attn_head_mask\"]\n",
    "        ):\n",
    "            if attn_mask is not None:\n",
    "                if attn_mask.size()[0] != (len(self.layers)):\n",
    "                    raise ValueError(\n",
    "                        f\"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for\"\n",
    "                        f\" {head_mask.size()[0]}.\"\n",
    "                    )\n",
    "\n",
    "        for iteration in range(num_iterations):\n",
    "            for idx, decoder_layer in enumerate(self.layers):\n",
    "                # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)\n",
    "                if output_hidden_states:\n",
    "                    all_hidden_states += (hidden_states,)\n",
    "                if self.training:\n",
    "                    dropout_probability = torch.rand([])\n",
    "                    if dropout_probability < self.layerdrop:\n",
    "                        continue\n",
    "\n",
    "                past_key_value = (\n",
    "                    past_key_values[idx] if past_key_values is not None else None\n",
    "                )\n",
    "\n",
    "                if self.gradient_checkpointing and self.training:\n",
    "                    layer_outputs = self._gradient_checkpointing_func(\n",
    "                        decoder_layer.__call__,\n",
    "                        hidden_states,\n",
    "                        attention_mask,\n",
    "                        encoder_hidden_states,\n",
    "                        encoder_attention_mask,\n",
    "                        head_mask[idx] if head_mask is not None else None,\n",
    "                        cross_attn_head_mask[idx]\n",
    "                        if cross_attn_head_mask is not None\n",
    "                        else None,\n",
    "                        None,\n",
    "                        output_attentions,\n",
    "                        use_cache,\n",
    "                    )\n",
    "                else:\n",
    "                    layer_outputs = decoder_layer(\n",
    "                        hidden_states,\n",
    "                        attention_mask=attention_mask,\n",
    "                        encoder_hidden_states=encoder_hidden_states,\n",
    "                        encoder_attention_mask=encoder_attention_mask,\n",
    "                        layer_head_mask=(\n",
    "                            head_mask[idx] if head_mask is not None else None\n",
    "                        ),\n",
    "                        cross_attn_layer_head_mask=(\n",
    "                            cross_attn_head_mask[idx]\n",
    "                            if cross_attn_head_mask is not None\n",
    "                            else None\n",
    "                        ),\n",
    "                        past_key_value=past_key_value,\n",
    "                        output_attentions=output_attentions,\n",
    "                        use_cache=use_cache,\n",
    "                    )\n",
    "                hidden_states = layer_outputs[0]\n",
    "\n",
    "                if use_cache:\n",
    "                    next_decoder_cache += (\n",
    "                        layer_outputs[3 if output_attentions else 1],\n",
    "                    )\n",
    "\n",
    "                if output_attentions:\n",
    "                    all_self_attns += (layer_outputs[1],)\n",
    "\n",
    "                    if encoder_hidden_states is not None:\n",
    "                        all_cross_attentions += (layer_outputs[2],)\n",
    "\n",
    "        # add hidden states from the last decoder layer\n",
    "        if output_hidden_states:\n",
    "            all_hidden_states += (hidden_states,)\n",
    "\n",
    "        next_cache = next_decoder_cache if use_cache else None\n",
    "        if not return_dict:\n",
    "            return tuple(\n",
    "                v\n",
    "                for v in [\n",
    "                    hidden_states,\n",
    "                    next_cache,\n",
    "                    all_hidden_states,\n",
    "                    all_self_attns,\n",
    "                    all_cross_attentions,\n",
    "                ]\n",
    "                if v is not None\n",
    "            )\n",
    "        return BaseModelOutputWithPastAndCrossAttentions(\n",
    "            last_hidden_state=hidden_states,\n",
    "            past_key_values=next_cache,\n",
    "            hidden_states=all_hidden_states,\n",
    "            attentions=all_self_attns,\n",
    "            cross_attentions=all_cross_attentions,\n",
    "        )  ## Decoder Layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Bart EncoderLayer Modification\n",
    "# class CustomBartEncoderLayer(nn.Module):\n",
    "#     def __init__(self, config: BartConfig, layer):\n",
    "#         super().__init__()\n",
    "#         self.layer = layer\n",
    "#         num_replicas = 19\n",
    "#         self.num_replicas = num_replicas\n",
    "#         self.embed_dim = config.d_model\n",
    "\n",
    "#         self.fc1 = nn.ModuleList(\n",
    "#             [nn.Linear(self.embed_dim, self.embed_dim) for _ in range(num_replicas)]\n",
    "#         )\n",
    "#         # print(self.embed_dim, config.decoder_ffn_dim)\n",
    "#         # self.fc2 = nn.ModuleList(\n",
    "#         #     [nn.Linear(self.embed_dim, self.embed_dim) for _ in range(num_replicas)]\n",
    "#         # )\n",
    "\n",
    "#         self.q1 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "#         self.k1 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "\n",
    "#         # self.q2 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "#         # self.k2 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "\n",
    "#         # print(self.embed_dim, config.decoder_ffn_dim, config.decoder_attention_heads)\n",
    "\n",
    "#         self.final_layer_norm = nn.LayerNorm(self.embed_dim)\n",
    "#         self.dropout = config.dropout\n",
    "#         self.activation_fn = ACT2FN[config.activation_function]\n",
    "#         self.activation_dropout = config.activation_dropout  # Define activation dropout\n",
    "\n",
    "#     def forward(self, x, *args, **kwargs):\n",
    "#         outputs = self.layer(x, *args, **kwargs)\n",
    "#         hidden_states = outputs[0]\n",
    "#         residual = hidden_states\n",
    "#         hidden_states = self.activation_fn(hidden_states)\n",
    "\n",
    "#         batch_size = hidden_states.size(1)\n",
    "\n",
    "#         q1 = self.q1(hidden_states)\n",
    "#         fc1_concat = torch.stack([fc1(hidden_states) for fc1 in self.fc1], dim=0)\n",
    "#         p1 = torch.stack([fc1.weight.data for fc1 in self.fc1], dim=0)\n",
    "#         p1 = p1.unsqueeze(2).expand(-1, -1, batch_size, -1)\n",
    "#         k1 = self.k1(p1)\n",
    "#         attention_weight_1 = scaled_dot_product_attention(q1, k1, self.embed_dim)\n",
    "#         attention_weight_1 = nn.functional.softmax(attention_weight_1, dim=0)\n",
    "#         # print(self.num_replicas)\n",
    "#         # print(\"q1\", q1.shape)\n",
    "#         # print(\"k1\", k1.shape)\n",
    "#         # print(\"fc1_concat\", fc1_concat.shape)\n",
    "#         # print(\"attention_weight_1\", attention_weight_1.shape)\n",
    "#         fc1_weighted = fc1_concat * attention_weight_1\n",
    "#         hidden_states = fc1_weighted.sum(dim=0)\n",
    "#         # hidden_states = self.activation_fn(hidden_states)\n",
    "#         # print((attention_weight_1.sum(dim=(1,2,3))))\n",
    "\n",
    "#         # hidden_states = nn.functional.dropout(\n",
    "#         #     hidden_states, p=self.activation_dropout, training=self.training\n",
    "#         # )\n",
    "\n",
    "#         # # Vectorized operation for fc2 layers\n",
    "#         # q2 = self.q2(hidden_states)\n",
    "#         # fc2_concat = torch.stack([fc2(hidden_states) for fc2 in self.fc2], dim=0)\n",
    "#         # p2 = torch.stack([fc2.weight.data for fc2 in self.fc2], dim=0)\n",
    "#         # p2 = p2.unsqueeze(2).expand(-1, -1, batch_size, -1)\n",
    "#         # k2 = self.k2(p2)\n",
    "#         # attention_weight_2 = scaled_dot_product_attention(q2, k2, self.embed_dim)\n",
    "#         # fc2_weighted = fc2_concat * attention_weight_2\n",
    "#         # hidden_states = fc2_weighted.mean(dim=0)\n",
    "\n",
    "#         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)\n",
    "\n",
    "#         hidden_states = hidden_states + residual\n",
    "#         hidden_states = self.final_layer_norm(hidden_states)\n",
    "\n",
    "#         return (hidden_states,) + outputs[1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BartDecoderLayer Modification\n",
    "class CustomBartDecoderLayer(nn.Module):\n",
    "    def __init__(self, config: BartConfig, layer):\n",
    "        super().__init__()\n",
    "        self.layer = layer\n",
    "        self.num_replicas = num_replicas\n",
    "        self.embed_dim = config.d_model  # Assuming embed_dim is d_model\n",
    "        self.fc1 = nn.ModuleList(\n",
    "            [\n",
    "                nn.Linear(self.embed_dim, self.embed_dim, bias=False)\n",
    "                for _ in range(num_replicas)\n",
    "            ]\n",
    "        )\n",
    "        # print(self.embed_dim, config.decoder_ffn_dim)\n",
    "        # self.fc2 = nn.ModuleList(\n",
    "        #     [nn.Linear(self.embed_dim, self.embed_dim, bias=False) for _ in range(num_replicas)]\n",
    "        # )\n",
    "        self.q1 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "        self.k1 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "\n",
    "        self.attention_weight_bias = nn.Parameter(torch.zeros(num_replicas, 1, 1, 1))\n",
    "\n",
    "        # self.q2 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "        # self.k2 = nn.Linear(self.embed_dim, self.embed_dim)\n",
    "\n",
    "        # print(self.embed_dim, config.decoder_ffn_dim, config.decoder_attention_heads)\n",
    "\n",
    "        self.final_layer_norm = nn.LayerNorm(self.embed_dim)\n",
    "        self.dropout = config.dropout\n",
    "        self.activation_fn = ACT2FN[config.activation_function]\n",
    "        self.activation_dropout = config.activation_dropout\n",
    "        self.mh_attn = VectorizedFeatureSpecificMultiHeadAttention(\n",
    "            num_heads=config.decoder_attention_heads, d_model=self.embed_dim, num_features=self.num_replicas\n",
    "        )\n",
    "        # config.decoder_attention_heads\n",
    "        # for i, fc1 in enumerate(self.fc1):\n",
    "        #     # Initialize with a simple pattern, e.g., all elements in the weight matrix are set to the index of the layer\n",
    "        #     nn.init.constant_(fc1.weight, 17 * i)\n",
    "\n",
    "    def forward(self, x, *args, **kwargs):\n",
    "        outputs = self.layer(x, *args, **kwargs)\n",
    "        hidden_states = outputs[0]\n",
    "        residual = hidden_states\n",
    "\n",
    "        batch_size = hidden_states.size(0)\n",
    "        seq_len = hidden_states.size(1)\n",
    "\n",
    "        p1 = torch.stack([fc1.weight for fc1 in self.fc1], dim=0)\n",
    "        # .transpose(0, 1)\n",
    "        p1 = p1.unsqueeze(1)\n",
    "        # .expand(-1, batch_size, -1, -1)\n",
    "\n",
    "        hidden_states_reshaped = hidden_states.unsqueeze(0).expand(\n",
    "            self.num_replicas, -1, -1, -1\n",
    "        )\n",
    "        fc1_concat = torch.matmul(hidden_states_reshaped, p1)\n",
    "        # fc1_concat1 = torch.stack([fc1(hidden_states) for fc1 in self.fc1], dim=0)\n",
    "\n",
    "        # print(\"p1\", p1.shape)\n",
    "        # print(\"hidden_states_reshaped\",hidden_states_reshaped.shape)\n",
    "        # print(\"fc1_concat1\",fc1_concat1.shape)\n",
    "        # print(\"fc1_concat\", fc1_concat.shape)\n",
    "        # torch.set_printoptions(sci_mode=False)\n",
    "        # print(\"Same:\", torch.allclose(fc1_concat, fc1_concat1, rtol = 0.001))\n",
    "        # print(torch.abs(fc1_concat - fc1_concat1).mean())\n",
    "        # print(torch.abs(fc1_concat - fc1_concat1).max())\n",
    "        # print()\n",
    "\n",
    "        # q1 = self.q1(hidden_states).unsqueeze(0).expand(self.num_replicas, -1, -1, -1)\n",
    "        # k1 = self.k1(p1)\n",
    "        # attention_weight_1 = scaled_dot_product_attention(q1, k1, self.embed_dim)        \n",
    "        # attention_weight_1 = attention_weight_1 + self.attention_weight_bias\n",
    "\n",
    "        # print(\"q1\", q1.shape)\n",
    "        # print(\"k1\", k1.shape)\n",
    "        mh_attn_weight = self.mh_attn(query=hidden_states_reshaped, key=p1) # shape: (num_experts, batch_size, num_heads, seq_len, embed_dim)\n",
    "        # print(\"attention_weight_1\", attention_weight_1.shape)\n",
    "        # print(\"mh_attn_weight\", mh_attn_weight.shape)\n",
    "        # print(\"Same:\", torch.allclose(attention_weight_1, mh_attn_weight, rtol=0.001))\n",
    "        # print(torch.abs(attention_weight_1 - mh_attn_weight).mean())\n",
    "        # print(torch.abs(attention_weight_1 - mh_attn_weight).max())\n",
    "        # print()\n",
    "\n",
    "        # attention_weight_1_norm_expert = nn.functional.softmax(\n",
    "        #     mh_attn_weight, dim=0\n",
    "        # ).mean(dim = 2)\n",
    "        # attention_weight_1_norm_feature = nn.functional.softmax(\n",
    "        #     mh_attn_weight, dim=-1\n",
    "        # ).mean(dim = 2)\n",
    "        # mh_attn_weight = (\n",
    "        #     3 * attention_weight_1_norm_expert + attention_weight_1_norm_feature\n",
    "        # )\n",
    "        mh_attn_weight = mh_attn_weight.mean(dim = 2)\n",
    "        # print(self.num_replicas)\n",
    "        # print(\"q1\", q1.shape)\n",
    "        # print(\"k1\", k1.shape)\n",
    "        # print(\"fc1_concat\", fc1_concat.shape)\n",
    "        # print(\"attention_weight_1\", attention_weight_1.shape)\n",
    "        fc1_weighted = fc1_concat * mh_attn_weight\n",
    "        hidden_states = (fc1_weighted.mean(dim=0))\n",
    "        # self.activation_fn\n",
    "        # print((attention_weight_1.sum(dim=(1,2,3))))\n",
    "        # hidden_states = nn.functional.dropout(\n",
    "        #     hidden_states, p=self.activation_dropout, training=self.training\n",
    "        # )\n",
    "\n",
    "        # # Vectorized operation for fc2 layers\n",
    "        # q2 = self.q2(hidden_states)\n",
    "        # fc2_concat = torch.stack([fc2(hidden_states) for fc2 in self.fc2], dim=0)\n",
    "        # p2 = torch.stack([fc2.weight.data for fc2 in self.fc2], dim=0)\n",
    "        # p2 = p2.unsqueeze(2).expand(-1, -1, batch_size, -1)\n",
    "        # k2 = self.k2(p2)\n",
    "        # attention_weight_2 = scaled_dot_product_attention(q2, k2, self.embed_dim)\n",
    "        # fc2_weighted = fc2_concat * attention_weight_2\n",
    "        # hidden_states = fc2_weighted.mean(dim=0)\n",
    "\n",
    "        hidden_states = nn.functional.dropout(\n",
    "            hidden_states, p=self.dropout, training=self.training\n",
    "        )\n",
    "\n",
    "        hidden_states = hidden_states + residual\n",
    "        hidden_states = self.final_layer_norm(hidden_states)\n",
    "\n",
    "        return (hidden_states,) + outputs[1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BartForConditionalGeneration(\n",
       "  (model): BartModel(\n",
       "    (shared): Embedding(50264, 1024, padding_idx=1)\n",
       "    (encoder): BartEncoder(\n",
       "      (embed_tokens): Embedding(50264, 1024, padding_idx=1)\n",
       "      (embed_positions): BartLearnedPositionalEmbedding(1026, 1024)\n",
       "      (layers): ModuleList(\n",
       "        (0-11): 12 x BartEncoderLayer(\n",
       "          (self_attn): BartAttention(\n",
       "            (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (v_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (q_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          )\n",
       "          (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (activation_fn): GELUActivation()\n",
       "          (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (layernorm_embedding): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (decoder): CustomBartDecoder(\n",
       "      (embed_tokens): Embedding(50264, 1024, padding_idx=1)\n",
       "      (embed_positions): BartLearnedPositionalEmbedding(1026, 1024)\n",
       "      (layers): ModuleList(\n",
       "        (0-11): 12 x CustomBartDecoderLayer(\n",
       "          (layer): BartDecoderLayer(\n",
       "            (self_attn): BartAttention(\n",
       "              (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (v_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (q_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            )\n",
       "            (activation_fn): GELUActivation()\n",
       "            (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (encoder_attn): BartAttention(\n",
       "              (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (v_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (q_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            )\n",
       "            (encoder_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          )\n",
       "          (fc1): ModuleList(\n",
       "            (0): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "          )\n",
       "          (q1): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (k1): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (activation_fn): GELUActivation()\n",
       "          (mh_attn): VectorizedFeatureSpecificMultiHeadAttention(\n",
       "            (wq): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (wk): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (layernorm_embedding): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (lm_head): Linear(in_features=1024, out_features=50264, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load the pretrained BART model\n",
    "model_name = \"facebook/bart-large-cnn\"\n",
    "config = BartConfig.from_pretrained(model_name)\n",
    "# print(config)\n",
    "model = BartForConditionalGeneration.from_pretrained(model_name, config=config)\n",
    "tokenizer = BartTokenizer.from_pretrained(model_name)\n",
    "\n",
    "\n",
    "num_iterations = 2\n",
    "num_replicas = 1\n",
    "model.model.decoder = CustomBartDecoder(\n",
    "    config=model.config, embed_tokens=model.model.shared\n",
    ")\n",
    "\n",
    "# Check if CUDA GPU is available\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "\n",
    "model_weights_path = f\".\\model_weights\\{model_name}.pth\"\n",
    "\n",
    "# # Load model weights to the device\n",
    "if torch.cuda.is_available():\n",
    "    model.load_state_dict(torch.load(model_weights_path, map_location=\"cuda\"))\n",
    "else:\n",
    "    model.load_state_dict(torch.load(model_weights_path, map_location=\"cpu\"))\n",
    "\n",
    "# Replace all customized layers\n",
    "for i, layer in enumerate(model.model.decoder.layers):\n",
    "    model.model.decoder.layers[i] = CustomBartDecoderLayer(model.config, layer)\n",
    "\n",
    "\n",
    "# Move the model to the specified device\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code Warehouse\n",
    "# model = BartForConditionalGeneration(config=config)\n",
    "\n",
    "# for i, layer in enumerate(model.model.encoder.layers):\n",
    "#     model.model.encoder.layers[i] = CustomBartEncoderLayer(model.config, layer)\n",
    "\n",
    "\n",
    "# #save pretrained model weights\n",
    "# torch.save(model.state_dict(), model_weigts_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of parameters: 469278732\n"
     ]
    }
   ],
   "source": [
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(\"Total number of parameters:\", total_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Freeze pretrained weight\n",
    "\n",
    "# Step 1: Freeze all pretrained weights\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "# Step 2: Unfreeze the weights in custom layers\n",
    "for i in range(len(model.model.decoder.layers)):\n",
    "    layer = model.model.decoder.layers[i]\n",
    "    if isinstance(layer, CustomBartDecoderLayer):\n",
    "        for param in layer.parameters():\n",
    "            param.requires_grad = True\n",
    "\n",
    "    if isinstance(layer, BartDecoderLayer):\n",
    "        for param in layer.parameters():\n",
    "            param.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference / Generate Text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NineProsecutorsLLProsecutorsLProsecutorsProsecutorsProsecutorsNineProsecutorsMarProsecutorsNineLNineProsecutorsLNineNineProsecutorsNineNineLProsecutorsNineAfterProsecutorsNineAtNineProsecutorsProsecutorsLLNineLLLProsecutorsAtProsecutorsProsecutorsAtLProsecutorsAfterProsecutorsLBarLNineWomanProsecutorsNineInProsecutorsProsecutorsFiveProsecutorsProsecutorsBarProsecutorsProsecutorsAfterNineProsecutorsInProsecutorsLAfterProsecutorsProsecutorsWomanProsecutorsProsecutorsPoliceProsecutorsNineMarNineNineNineAtProsecutorsLInProsecutorsWomanNineProsecutorsAtNineNineAfterLLAtLLInLNineInNineNineWomanLProsecutorsBarNineProsecutorsThereProsecutorsNineBarLProsecutorsInLAtInNineProsecutorsAfterLProsecutorsWomanLLAroundProsecutorsProsecutors\n"
     ]
    }
   ],
   "source": [
    "# # Sample text to summarize\n",
    "text = \"\"\"New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County, New York.\n",
    "A year later, she got married again in Westchester County, but to a different man and without divorcing her first husband.\n",
    "Only 18 days after that marriage, she got hitched yet again. Then, Barrientos declared \"I do\" five more times, sometimes only within two weeks of each other.\n",
    "In 2010, she married once more, this time in the Bronx. In an application for a marriage license, she stated it was her \"first and only\" marriage.\n",
    "Barrientos, now 39, is facing two criminal counts of \"offering a false instrument for filing in the first degree,\" referring to her false statements on the\n",
    "2010 marriage license application, according to court documents.\n",
    "Prosecutors said the marriages were part of an immigration scam.\n",
    "On Friday, she pleaded not guilty at State Supreme Court in the Bronx, according to her attorney, Christopher Wright, who declined to comment further.\n",
    "After leaving court, Barrientos was arrested and charged with theft of service and criminal trespass for allegedly sneaking into the New York subway through an emergency exit, said Detective\n",
    "Annette Markowski, a police spokeswoman. In total, Barrientos has been married 10 times, with nine of her marriages occurring between 1999 and 2002.\n",
    "All occurred either in Westchester County, Long Island, New Jersey or the Bronx. She is believed to still be married to four men, and at one time, she was married to eight men at once, prosecutors say.\n",
    "Prosecutors said the immigration scam involved some of her husbands, who filed for permanent residence status shortly after the marriages.\n",
    "Any divorces happened only after such filings were approved. It was unclear whether any of the men will be prosecuted.\n",
    "The case was referred to the Bronx District Attorney\\'s Office by Immigration and Customs Enforcement and the Department of Homeland Security\\'s\n",
    "Investigation Division. Seven of the men are from so-called \"red-flagged\" countries, including Egypt, Turkey, Georgia, Pakistan and Mali.\n",
    "Her eighth husband, Rashid Rajput, was deported in 2006 to his native Pakistan after an investigation by the Joint Terrorism Task Force.\n",
    "If convicted, Barrientos faces up to four years in prison.  Her next court appearance is scheduled for May 18.\"\"\"\n",
    "\n",
    "# Encode the text into tokens\n",
    "inputs = tokenizer([text], return_tensors=\"pt\")\n",
    "# , max_length=1024\n",
    "\n",
    "# Move the input tensors to the same device as the model\n",
    "inputs = inputs.to(device)\n",
    "\n",
    "# Generate a summary of the encoded text\n",
    "summary_ids = model.generate(\n",
    "    inputs[\"input_ids\"],\n",
    "    num_beams=4,\n",
    "    # max_length=51,\n",
    "    # early_stopping=True\n",
    ")\n",
    "\n",
    "# Decode the summary\n",
    "summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)\n",
    "print(summary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " mitochond mitochond mitochond mitochond tsun mitochond mitochond tsun mitochond tsun mitochond tsun metaphors mitochond mitochond metaphors mitochond metaphors mitochond metaphors cannibal mitochond mitochond cannibal mitochond cannibal mitochond cannibal hypothal mitochond mitochond hypothal mitochond hypothal mitochond hypothal starve mitochond mitochond starve mitochond starve mitochond starve eaten mitochond mitochond eaten mitochond eaten mitochond eaten disag mitochond mitochond disag mitochond disag mitochond disag Apostles mitochond mitochond Apostles mitochond Apostles mitochond Apostles Dalai mitochond mitochond Dalai mitochond Dalai mitochond Dalai resemb mitochond mitochond resemb mitochond resemb mitochond resemb apostles mitochond mitochond apostles mitochond apostles mitochond apostles\n"
     ]
    }
   ],
   "source": [
    "# Prepare the batched input\n",
    "input_texts = [\n",
    "    \"Who is the president of China?\",\n",
    "    \"Who is the president of the US?\",\n",
    "    \"Who is the president of the Russia?\",\n",
    "    # \"\"\"New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County, New York. A year later, she got married again in Westchester County, but to a different man and without divorcing her first husband. Only 18 days after that marriage, she got hitched yet again. Then, Barrientos declared \"I do\" five more times, sometimes only within two weeks of each other. In 2010, she married once more, this time in the Bronx. In an application for a marriage license, she stated it was her \"first and only\" marriage.\"\"\",\n",
    "    # \"\"\"Your model's primary bottlenecks appear to be matrix multiplication and linear layer operations, both in terms of computation and possibly memory usage. Focusing your optimization efforts on these areas, along with minimizing unnecessary memory operations, could lead to significant improvements in performance. Remember, optimizations can sometimes affect model accuracy, so it's important to validate your model's performance after making any changes.\"\"\",\n",
    "    \"\"\"Which technology was developed most recently? Options:\n",
    "    A. television    \n",
    "    B. refrigerator    \n",
    "    C. cellular telephone\n",
    "    D. airplane\n",
    "    Which is the correct answer?\"\"\",\n",
    "]\n",
    "\n",
    "input_texts = \"I eat\"\n",
    "\n",
    "inputs = tokenizer(\n",
    "    input_texts,\n",
    "    return_tensors=\"pt\",\n",
    "    padding=True,\n",
    "    # truncation=True\n",
    ")\n",
    "\n",
    "# Move the input tensors to the same device as the model\n",
    "inputs = inputs.to(device)\n",
    "\n",
    "# Generate the output\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    output_tokens = model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        attention_mask=inputs[\"attention_mask\"],\n",
    "        num_beams=4,\n",
    "        # max_length=50,  # Optionally set a max length if desired\n",
    "    )\n",
    "\n",
    "# Decode the generated tokens for each input in the batch\n",
    "output_texts = [\n",
    "    tokenizer.decode(token, skip_special_tokens=True) for token in output_tokens\n",
    "]\n",
    "\n",
    "# Print each output\n",
    "for output in output_texts:\n",
    "    print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torch.profiler import profile, record_function, ProfilerActivity\n",
    "\n",
    "# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],\n",
    "#              record_shapes=True) as prof:\n",
    "#     with record_function(\"model_inference\"):\n",
    "#         # Your model inference code here\n",
    "#         outputs = model.generate(\n",
    "#         inputs[\"input_ids\"],\n",
    "#         attention_mask=inputs[\"attention_mask\"],\n",
    "#         num_beams=4,\n",
    "#         # max_length=50,  # Optionally set a max length if desired\n",
    "#         )\n",
    "#         # model(input_data)\n",
    "\n",
    "# print(prof.key_averages().table(sort_by=\"cuda_time_total\", row_limit=10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training BART"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\stbyu\\anaconda3\\envs\\eecs595python10\\lib\\site-packages\\transformers\\optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "  0%|          | 0/1 [00:00<?, ?it/s]C:\\Users\\stbyu\\AppData\\Local\\Temp\\ipykernel_18956\\3573005218.py:10: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  return {key: torch.tensor(val).squeeze() for key, val in self.encodings[idx].items()}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input IDs shape: torch.Size([2, 1024])\n",
      "Attention mask shape: torch.Size([2, 1024])\n",
      "Input IDs shape: torch.Size([1, 1024])\n",
      "Attention mask shape: torch.Size([1, 1024])\n",
      "epoch 1 loss "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:43<00:00, 43.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21.4877, device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "texts = [\"I eat apple\", \"text2...\", \"text3...\"]  # Replace with your corpus\n",
    "\n",
    "tokenized_texts = [tokenizer(text, return_tensors=\"pt\", max_length=1024, truncation=True, padding=\"max_length\") for text in texts]\n",
    "\n",
    "class TextDataset(Dataset):\n",
    "    def __init__(self, encodings):\n",
    "        self.encodings = encodings\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {key: torch.tensor(val).squeeze() for key, val in self.encodings[idx].items()}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.encodings)\n",
    "\n",
    "\n",
    "dataset = TextDataset(tokenized_texts)\n",
    "dataloader = DataLoader(dataset, batch_size=2, shuffle=True)\n",
    "\n",
    "## Training Loop\n",
    "\n",
    "optimizer = AdamW(model.parameters(), lr=5e-5)\n",
    "\n",
    "num_epochs = 1\n",
    "model.train()\n",
    "for epoch in tqdm(range(num_epochs)):  # Define num_epochs\n",
    "    epoch_loss = 0\n",
    "    for batch in dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        attention_mask = batch['attention_mask'].to(device)\n",
    "        \n",
    "        # Debugging: Print shapes\n",
    "        print(f\"Input IDs shape: {input_ids.shape}\")\n",
    "        print(f\"Attention mask shape: {attention_mask.shape}\")\n",
    "        labels = input_ids.clone()  # Assuming a denoising objective\n",
    "        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "        loss = outputs.loss\n",
    "        epoch_loss += loss\n",
    "        loss.backward()\n",
    "        optimizer.step()    \n",
    "    epoch_loss /= len(texts)\n",
    "    print(f\"epoch {epoch+1} loss\", epoch_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"path/to/save/model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "eecs595python10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}