{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "from transformers import FlaxRobertaModel, RobertaTokenizerFast\n",
    "from datasets import load_dataset\n",
    "import jax\n",
    "\n",
    "dataset = load_dataset('oscar', \"unshuffled_deduplicated_en\", split='train', streaming=True)\n",
    "\n",
    "dummy_input = next(iter(dataset))[\"text\"]\n",
    "\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"roberta-base\")\n",
    "input_ids = tokenizer(dummy_input, return_tensors=\"np\").input_ids[:, :10]\n",
    "\n",
    "model = FlaxRobertaModel.from_pretrained(\"julien-c/dummy-unknown\")\n",
    "\n",
    "# run a forward pass, should return an object `FlaxBaseModelOutputWithPooling`\n",
    "model(input_ids)"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/5.58k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e82d4ce42bd44b3da2b9f7099b1d41b6",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/359k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aaff74f2a67b402e8df4920cf17dc897",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "615c025708384873b750fcce8656821b",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ec2e1639ea924d0ea34334934676fd4c",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a76e34bf8131437aa1913e05773a2f7a",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (942 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/496 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3935e71c772c41d0a90ae657716848d8",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading:   0%|          | 0.00/58.5k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb26132349f54c378b2ae018ae95c658",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "INFO:absl:Starting the local TPU driver.\n",
      "INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://\n",
      "INFO:absl:Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: \"cuda\". Available platform names are: Interpreter Host TPU\n",
      "Some weights of FlaxRobertaModel were not initialized from the model checkpoint at julien-c/dummy-unknown and are newly initialized: {('pooler', 'dense', 'bias'), ('pooler', 'dense', 'kernel')}\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "FlaxBaseModelOutputWithPooling(last_hidden_state=DeviceArray([[[-4.73004431e-01, -1.12082589e+00, -1.46316242e+00,\n",
       "                9.30495918e-01, -3.00514668e-01,  8.28465223e-01,\n",
       "               -5.99549664e-03, -2.08820820e+00,  5.73224485e-01,\n",
       "                1.13552138e-01,  9.50309098e-01,  1.82359887e-04,\n",
       "               -2.36949772e-01,  3.99159789e-02, -1.58831275e+00,\n",
       "                1.21534312e+00,  1.61006236e+00, -6.29136503e-01,\n",
       "                1.42789042e+00,  2.16668591e-01],\n",
       "              [ 1.05048262e-01,  4.38308679e-02,  7.17148423e-01,\n",
       "               -8.70970964e-01, -9.46345925e-02,  1.24356456e-01,\n",
       "               -1.38944432e-01,  9.06061828e-01, -5.45630276e-01,\n",
       "                8.97509232e-02, -5.71692824e-01,  1.93972692e-01,\n",
       "               -1.39611530e+00, -2.21010113e+00, -5.64776003e-01,\n",
       "                4.50222403e-01,  1.92772031e+00, -1.16361332e+00,\n",
       "                1.13953459e+00,  1.85883164e+00],\n",
       "              [-1.38152033e-01, -1.94758683e-01,  6.30050123e-01,\n",
       "               -5.49758852e-01, -9.73834991e-01, -4.73883361e-01,\n",
       "                1.52569726e-01,  1.14790475e+00, -1.40914500e-01,\n",
       "               -6.94071293e-01, -2.86773652e-01,  1.09020483e+00,\n",
       "               -1.16347635e+00, -2.10751987e+00, -5.70781708e-01,\n",
       "                9.34546471e-01,  1.95043743e+00, -5.05802572e-01,\n",
       "               -7.99659565e-02,  1.97398031e+00],\n",
       "              [ 4.59462494e-01,  1.81284267e-02,  3.17216758e-03,\n",
       "               -1.24535978e+00, -5.23790121e-01, -6.02854550e-01,\n",
       "               -1.39238501e+00, -4.20255251e-02, -9.87934768e-01,\n",
       "               -9.76055562e-02, -4.30685580e-01,  1.21701121e+00,\n",
       "               -5.81501126e-02, -6.56547129e-01, -2.54513383e-01,\n",
       "                8.32962990e-01,  2.22358274e+00, -8.06928217e-01,\n",
       "               -1.45050630e-01,  2.48951077e+00],\n",
       "              [ 9.24803019e-02,  4.79857206e-01,  1.03523517e+00,\n",
       "               -9.18277562e-01, -9.84333336e-01, -2.20230371e-01,\n",
       "               -5.54836631e-01,  1.23434842e+00,  6.11479282e-01,\n",
       "                9.24420595e-01, -4.94961560e-01,  7.32724786e-01,\n",
       "               -1.68369830e+00, -1.15283561e+00, -1.46726906e+00,\n",
       "                2.32458755e-01,  2.02297497e+00, -6.97842836e-01,\n",
       "               -5.20550787e-01,  1.32885659e+00],\n",
       "              [ 6.31922722e-01, -1.04144073e+00, -1.32531202e+00,\n",
       "                1.33779496e-01,  5.81125140e-01, -3.91236767e-02,\n",
       "                7.17227280e-01,  6.58484936e-01, -6.44005299e-01,\n",
       "                3.44709679e-02, -6.01282239e-01,  9.57775652e-01,\n",
       "               -8.74027431e-01, -1.37680268e+00, -1.04416537e+00,\n",
       "                1.03142655e+00,  2.40814567e+00, -1.29077399e+00,\n",
       "               -1.27586871e-01,  1.21016145e+00],\n",
       "              [ 2.51221180e-01, -8.25482070e-01,  6.45016313e-01,\n",
       "               -1.90095508e+00, -7.89572060e-01, -2.49178708e-01,\n",
       "                6.08372211e-01,  6.67878211e-01, -4.74199533e-01,\n",
       "               -5.43338001e-01,  5.44238031e-01,  8.53681684e-01,\n",
       "               -3.33928257e-01, -1.40273571e+00, -4.08071905e-01,\n",
       "                1.68641591e+00,  1.90652800e+00, -8.37633669e-01,\n",
       "               -8.33109736e-01,  1.43485296e+00],\n",
       "              [ 5.63060522e-01, -1.37411666e+00,  1.77591711e-01,\n",
       "               -1.41981161e+00,  2.86106497e-01, -9.85570014e-01,\n",
       "               -1.02347910e+00,  1.08637953e+00, -6.72662184e-02,\n",
       "                5.75017147e-02, -3.98549348e-01,  9.95321929e-01,\n",
       "               -8.35580587e-01, -1.78589416e+00,  1.59830645e-01,\n",
       "                1.17250645e+00,  1.52774048e+00, -1.73241213e-01,\n",
       "                1.29489526e-01,  1.90797997e+00],\n",
       "              [-4.75205094e-01,  5.75383492e-02, -6.28697217e-01,\n",
       "               -7.07582593e-01,  4.09776330e-01, -1.17534667e-01,\n",
       "                3.11417371e-01,  1.34317398e+00, -4.73465234e-01,\n",
       "               -2.62839556e-01, -9.32353258e-01,  1.41049230e+00,\n",
       "               -1.52682090e+00, -8.14880610e-01, -8.40299845e-01,\n",
       "                1.79613554e+00,  9.69104469e-01, -1.18959296e+00,\n",
       "               -3.80980015e-01,  2.05261350e+00],\n",
       "              [-1.20868230e+00, -1.60126913e+00,  5.79116225e-01,\n",
       "               -7.73360193e-01,  7.20766857e-02, -2.48067677e-01,\n",
       "               -9.75876153e-01,  1.74336642e-01, -1.48872507e+00,\n",
       "               -5.56367226e-02, -1.44585848e-01,  1.34230399e+00,\n",
       "                5.55252552e-01, -6.05108261e-01,  1.63732314e+00,\n",
       "                1.46649110e+00,  1.96676946e+00, -2.72412688e-01,\n",
       "                2.44762152e-01, -6.64708257e-01]]], dtype=float32), pooler_output=DeviceArray([[-0.04240899,  0.01514644,  0.16116808, -0.1028635 ,\n",
       "               0.08844209, -0.02510494,  0.12796836, -0.0424355 ,\n",
       "               0.06956789,  0.07822348, -0.01671289,  0.00228395,\n",
       "               0.07130629,  0.10654891,  0.10498928,  0.11156421,\n",
       "              -0.01685666, -0.10740762, -0.07814675,  0.12726422]],            dtype=float32), hidden_states=None, attentions=None)"
      ]
     },
     "metadata": {},
     "execution_count": 1
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import jax\n",
    "jax.device_count()"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "8"
      ]
     },
     "metadata": {},
     "execution_count": 1
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "import torch_xla\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "import librosa"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "import numpy as np"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "(1-np.array([[True, False]]))*-100"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "array([[   0, -100]])"
      ]
     },
     "metadata": {},
     "execution_count": 7
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e72d7b5d53f1ae10e392d8735559d980265e0a06ad4750ac62c5f8e194e9295d"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit ('env_wav')",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}