{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ddb479e-9d7e-4392-8fc0-fd1c66a07a2b",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['bert.pooler.dense.bias', 'bert.pooler.dense.weight', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 512])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import transformers\n",
    "transformers.set_seed(42)\n",
    "device = \"cuda\"\n",
    "# import sys\n",
    "\n",
    "# for key in list(sys.modules):\n",
    "#     if key.startswith(\"transformers.\"):\n",
    "#         sys.modules.pop(key)\n",
    "\n",
    "from transformers import AutoModelForMaskedLM\n",
    "model = AutoModelForMaskedLM.from_pretrained(\"bert-base-uncased\", torch_dtype=torch.float16, attn_implementation=\"sdpa\").to(device)\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "\n",
    "model.config.alek_says_ltr = True\n",
    "model.config.alek_says_rtl = False\n",
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset(\"Salesforce/wikitext\", \"wikitext-103-v1\")\n",
    "train_ds = ds[\"train\"]\n",
    "inputs = tokenizer(train_ds[10][\"text\"], return_tensors=\"pt\", padding='max_length', truncation=True)\n",
    "\n",
    "print(inputs[\"input_ids\"].size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "30472f6c-31d6-4768-8dca-d3535be28501",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "output = model(**{k: v.to(device) for k, v in inputs.items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ac46bf41-8cd3-4190-aa4b-6142d6d4d986",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ -7.1094,  -7.1445,  -7.2148,  ...,  -6.6484,  -7.0703,  -3.6758],\n",
       "         [-14.2969, -14.2656, -14.3828,  ..., -11.9766, -11.3281,  -9.4922],\n",
       "         [-10.7344, -10.6250, -10.7266,  ...,  -8.6641,  -8.2188,  -5.0859],\n",
       "         ...,\n",
       "         [ -3.4277,  -3.5664,  -3.9434,  ...,  -2.0000,  -4.4727,  -3.7148],\n",
       "         [ -3.7227,  -3.8770,  -4.2383,  ...,  -2.1367,  -4.5977,  -3.9336],\n",
       "         [ -4.2070,  -4.3672,  -4.7578,  ...,  -2.4941,  -4.7734,  -4.7227]]],\n",
       "       device='cuda:0', dtype=torch.float16, grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "baa757ff-3ba2-4a72-b819-a2283b729c18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-6.3281, -6.3555, -6.4531,  ..., -5.5234, -4.1797, -5.7891],\n",
       "         [-6.7891, -6.6914, -6.7812,  ..., -6.1680, -5.1094, -5.5273],\n",
       "         [-7.1641, -7.1055, -7.0625,  ..., -6.2383, -5.3711, -5.5273],\n",
       "         ...,\n",
       "         [ 9.4844,  8.9219,  9.2422,  ...,  7.6133,  7.2578,  9.9062],\n",
       "         [10.3672,  9.8516, 10.1797,  ...,  8.5547,  8.0781, 10.5938],\n",
       "         [ 8.3828,  8.0781,  8.1641,  ...,  7.2422,  6.7734,  7.9961]]],\n",
       "       device='cuda:0', dtype=torch.float16, grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a33632fb-ad41-49e3-acee-91d1dda974b8",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# output1 = model(**{k: v.to(device) for k, v in inputs.items()})\n",
    "# print(output1.logits)\n",
    "# output2 = model(**{k: v.to(device) for k, v in inputs.items()}, encoder_attention_mask=torch.zeros(1, 512, 512))\n",
    "# print(output2.logits)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}