{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup & Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting requirements.txt\n"
     ]
    }
   ],
   "source": [
    "%%writefile requirements.txt\n",
    "git+https://github.com/openai/whisper.git@8cf36f3508c9acd341a45eb2364239a3d81458b9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Test model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2022-09-23 20:32:18--  https://cdn-media.huggingface.co/speech_samples/sample1.flac\n",
      "Resolving cdn-media.huggingface.co (cdn-media.huggingface.co)... 13.32.151.62, 13.32.151.23, 13.32.151.60, ...\n",
      "Connecting to cdn-media.huggingface.co (cdn-media.huggingface.co)|13.32.151.62|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 282378 (276K) [audio/flac]\n",
      "Saving to: ‘sample1.flac’\n",
      "\n",
      "sample1.flac        100%[===================>] 275.76K  --.-KB/s    in 0.003s  \n",
      "\n",
      "2022-09-23 20:32:18 (78.7 MB/s) - ‘sample1.flac’ saved [282378/282378]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://cdn-media.huggingface.co/speech_samples/sample1.flac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████| 2.87G/2.87G [01:11<00:00, 42.9MiB/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Detected language: english\n",
      " going along slushy country roads and speaking to damp audiences in drafty school rooms day after day for a fortnight. he'll have to put in an appearance at some place of worship on sunday morning and he can come to us immediately afterwards.\n"
     ]
    }
   ],
   "source": [
    "import whisper\n",
    "\n",
    "model = whisper.load_model(\"large\")\n",
    "result = model.transcribe(\"sample1.flac\")\n",
    "print(result[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create Custom Handler for Inference Endpoints\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting handler.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile handler.py\n",
    "from typing import  Dict\n",
    "from transformers.pipelines.audio_utils import ffmpeg_read\n",
    "import whisper\n",
    "import torch\n",
    "\n",
    "SAMPLE_RATE = 16000\n",
    "\n",
    "\n",
    "\n",
    "class EndpointHandler():\n",
    "    def __init__(self, path=\"\"):\n",
    "        # load the model\n",
    "        self.model = whisper.load_model(\"medium\")\n",
    "\n",
    "\n",
    "    def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            data (:obj:):\n",
    "                includes the deserialized audio file as bytes\n",
    "        Return:\n",
    "            A :obj:`dict`:. base64 encoded image\n",
    "        \"\"\"\n",
    "        # process input\n",
    "        inputs = data.pop(\"inputs\", data)\n",
    "        audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)\n",
    "        audio_tensor= torch.from_numpy(audio_nparray)\n",
    "        \n",
    "        # run inference pipeline\n",
    "        result = self.model.transcribe(audio_nparray)\n",
    "\n",
    "        # postprocess the prediction\n",
    "        return {\"text\": result[\"text\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "test custom pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from handler import EndpointHandler\n",
    "\n",
    "# init handler\n",
    "my_handler = EndpointHandler(path=\".\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/endpoints/openai-whisper-endpoint/handler.py:27: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  ../torch/csrc/utils/tensor_numpy.cpp:178.)\n",
      "  audio_tensor= torch.from_numpy(audio_nparray)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Detected language: english\n"
     ]
    }
   ],
   "source": [
    "import base64\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "import json\n",
    "\n",
    "# file reader\n",
    "with open(\"sample1.flac\", \"rb\") as f:\n",
    "  request = {\"inputs\": f.read()}\n",
    "\n",
    "\n",
    "# test the handler\n",
    "pred = my_handler(request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'transcription': \" going along slushy country roads and speaking to damp audiences in draughty school rooms day after day for a fortnight. He'll have to put in an appearance at some place of worship on Sunday morning, and he can come to us immediately afterwards.\"}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'{\"transcription\": \" going along slushy country roads and speaking to damp audiences in draughty school rooms day after day for a fortnight. He\\'ll have to put in an appearance at some place of worship on Sunday morning, and he can come to us immediately afterwards.\"}'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "json.dumps({'transcription': \" going along slushy country roads and speaking to damp audiences in draughty school rooms day after day for a fortnight. He'll have to put in an appearance at some place of worship on Sunday morning, and he can come to us immediately afterwards.\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('dev': conda)",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f6dd96c16031089903d5a31ec148b80aeb0d39c32affb1a1080393235fbfa2fc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}