{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create an inference function\n",
    "\n",
    "Everything is ready to wrap the previously trained model and deploy it. \n",
    "\n",
    "First, we need to configure the LargeFileBackend, the TracingDatabase and GreatAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;226mThe value of `ENVIRONMENT` contains the \"ENV` prefix but `ENVIRONMENT` is not defined as an environment variable, using the default value defined above (`DEVELOPMENT`)\u001b[0m\n",
      "\u001b[38;5;226mEnvironment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
      "\u001b[38;5;39mMongoDbDriver has been already configured: skipping initialisation\u001b[0m\n",
      "\u001b[38;5;39mLargeFileS3 has been already configured: skipping initialisation\u001b[0m\n",
      "\u001b[38;5;39mGreatAI (v0.1.6): configured ✅\u001b[0m\n",
      "\u001b[38;5;39m  🔩 tracing_database: MongoDbDriver\u001b[0m\n",
      "\u001b[38;5;39m  🔩 large_file_implementation: LargeFileS3\u001b[0m\n",
      "\u001b[38;5;39m  🔩 is_production: False\u001b[0m\n",
      "\u001b[38;5;39m  🔩 should_log_exception_stack: True\u001b[0m\n",
      "\u001b[38;5;39m  🔩 prediction_cache_size: 4096\u001b[0m\n",
      "\u001b[38;5;39m  🔩 dashboard_table_size: 100\u001b[0m\n",
      "\u001b[38;5;226mYou still need to check whether you follow all best practices before trusting your deployment.\u001b[0m\n",
      "\u001b[38;5;226m> Find out more at https://se-ml.github.io/practices\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from great_ai.utilities import ConfigFile\n",
    "from great_ai.large_file import LargeFileS3\n",
    "from great_ai import configure, MongoDbDriver\n",
    "\n",
    "configuration = ConfigFile(\"config.ini\")\n",
    "\n",
    "LargeFileS3.configure_credentials_from_file(configuration)\n",
    "MongoDbDriver.configure_credentials_from_file(configuration)\n",
    "\n",
    "configure(\n",
    "    dashboard_table_size=100,  # traces are small, we can show many\n",
    "    prediction_cache_size=4096,  # predictions are expensive, cache them\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a pleasant developer experience, we create some typed models that will show up in the automatically generated OpenAPI schema specification and will also provide runtime type validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "class Attention(BaseModel):\n",
    "    weight: float\n",
    "    token: str\n",
    "\n",
    "\n",
    "class EvaluatedSentence(BaseModel):\n",
    "    score: float\n",
    "    text: str\n",
    "    explanation: List[Attention]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even though `@use_model` caches the remote files locally and it also handles deserialising objects, we only use it to store a directory. In this case, it gives back a path, the path to that directory. So, we need to load the files from that folder ourselves. In order to only load it once per process, we create a small model loader helper function.\n",
    "\n",
    "> This is usually not needed, however, when we can outsmart `dill` so for optimisation purposes, we do it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;39mLatest version of scibert-highlights is 0 (from versions: 0)\u001b[0m\n",
      "\u001b[38;5;39mFile scibert-highlights-0 found in cache\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from great_ai import use_model\n",
    "from pathlib import Path\n",
    "from typing import Tuple\n",
    "from transformers import (\n",
    "    PreTrainedModel,\n",
    "    PreTrainedTokenizer,\n",
    ")\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    ")\n",
    "\n",
    "_tokenizer: PreTrainedTokenizer = None\n",
    "_loaded_model: PreTrainedModel = None\n",
    "\n",
    "\n",
    "@use_model(\"scibert-highlights\", version=\"latest\", model_kwarg_name=\"model_path\")\n",
    "def get_tokenizer_and_model(\n",
    "    model_path: Path, original_model: str = \"allenai/scibert_scivocab_uncased\"\n",
    ") -> Tuple[PreTrainedTokenizer, PreTrainedModel]:\n",
    "    global _tokenizer, _loaded_model\n",
    "\n",
    "    if _tokenizer is None:\n",
    "        _tokenizer = AutoTokenizer.from_pretrained(original_model)\n",
    "\n",
    "    if _loaded_model is None:\n",
    "        config = AutoConfig.from_pretrained(\n",
    "            model_path, output_hidden_states=True, output_attentions=True\n",
    "        )\n",
    "        _loaded_model = AutoModelForSequenceClassification.from_pretrained(\n",
    "            model_path, config=config\n",
    "        )\n",
    "\n",
    "    return _tokenizer, _loaded_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, implement the inference function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from great_ai import GreatAI\n",
    "from great_ai.utilities import clean\n",
    "\n",
    "import re\n",
    "import numpy as np\n",
    "import torch\n",
    "from transformers.modeling_outputs import SequenceClassifierOutput\n",
    "\n",
    "\n",
    "@GreatAI.create\n",
    "def find_highlights(sentence: str) -> EvaluatedSentence:\n",
    "    \"\"\"Get the interestingness prediction of the input sentence using SciBERT.\n",
    "\n",
    "    Run the SciBERT model in inference mode and evaluate the sentence.\n",
    "    Additionally, provide explanation in the form of the last layer's sum attention\n",
    "    between `[CLS]` and the other tokens.\n",
    "    \"\"\"\n",
    "\n",
    "    tokenizer, loaded_model = get_tokenizer_and_model()\n",
    "    sentence = clean(sentence, convert_to_ascii=True, remove_brackets=True)\n",
    "\n",
    "    tensors = tokenizer(sentence, return_tensors=\"pt\", truncation=True, max_length=512)\n",
    "\n",
    "    with torch.inference_mode():\n",
    "        result: SequenceClassifierOutput = loaded_model(**tensors)\n",
    "        positive_likelihood = torch.nn.Softmax(dim=1)(result.logits)[0][1]\n",
    "    tokens = tensors[\"input_ids\"][0]\n",
    "\n",
    "    attentions = np.sum(result.attentions[-1].numpy()[0], axis=0)[0][1:-1]\n",
    "    # Tuple of `torch.FloatTensor` (one for each layer) of shape\n",
    "    # `(batch_size, num_heads, sequence_length, sequence_length)`.\n",
    "\n",
    "    explanation = []\n",
    "\n",
    "    token_attentions = list(zip(attentions, tokens[1:-1]))\n",
    "    for token in re.split(r\"([ .,])\", sentence):\n",
    "        token = token.strip()\n",
    "        if not token:\n",
    "            continue\n",
    "        bert_tokens = tokenizer(\n",
    "            token, return_tensors=\"pt\", truncation=True, max_length=512\n",
    "        )[\"input_ids\"][0][\n",
    "            1:-1\n",
    "        ]  # truncation=True needed to fix `RuntimeError: Already borrowed`\n",
    "        weight = 0\n",
    "        for t1 in bert_tokens:\n",
    "            if not token_attentions:\n",
    "                break\n",
    "            a, t2 = token_attentions.pop(0)\n",
    "            assert t1 == t2, sentence\n",
    "            weight += a\n",
    "        explanation.append(\n",
    "            Attention(\n",
    "                token=token if token in \".,\" else \" \" + token, weight=round(weight, 4)\n",
    "            )\n",
    "        )\n",
    "        if not token_attentions:\n",
    "            break\n",
    "\n",
    "    return EvaluatedSentence(\n",
    "        score=positive_likelihood, text=sentence, explanation=explanation\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A simple test to see everything works. Note that the models list is filled by the `@use_model` call even though it's not on the main inference function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Trace[EvaluatedSentence]({'created': '2022-07-16T18:47:29.581701',\n",
       "   'exception': None,\n",
       "   'feedback': None,\n",
       "   'logged_values': { 'arg:sentence:length': 51,\n",
       "                      'arg:sentence:value': 'Our solution has outperformed the '\n",
       "                                            'state-of-the-art.'},\n",
       "   'models': [{'key': 'scibert-highlights', 'version': 0}],\n",
       "   'original_execution_time_ms': 7127.2063,\n",
       "   'output': { 'explanation': [ {'token': ' Our', 'weight': 0.3993},\n",
       "                                {'token': ' solution', 'weight': 0.3481},\n",
       "                                {'token': ' has', 'weight': 0.2945},\n",
       "                                {'token': ' outperformed', 'weight': 0.4011},\n",
       "                                {'token': ' the', 'weight': 0.1484},\n",
       "                                {'token': ' state-of-the-art', 'weight': 0.5727},\n",
       "                                {'token': '.', 'weight': 7.775}],\n",
       "               'score': 0.9991180300712585,\n",
       "               'text': 'Our solution has outperformed the state-of-the-art.'},\n",
       "   'tags': ['find_highlights', 'online', 'development'],\n",
       "   'trace_id': '56e20e94-79df-4793-ae61-d20820ebe2d3'}),\n",
       " Trace[EvaluatedSentence]({'created': '2022-07-16T18:47:37.020275',\n",
       "   'exception': None,\n",
       "   'feedback': None,\n",
       "   'logged_values': { 'arg:sentence:length': 36,\n",
       "                      'arg:sentence:value': 'Their solution did not perform '\n",
       "                                            'well.'},\n",
       "   'models': [{'key': 'scibert-highlights', 'version': 0}],\n",
       "   'original_execution_time_ms': 170.7057,\n",
       "   'output': { 'explanation': [ {'token': ' Their', 'weight': 1.1475},\n",
       "                                {'token': ' solution', 'weight': 0.8205},\n",
       "                                {'token': ' did', 'weight': 0.3254},\n",
       "                                {'token': ' not', 'weight': 0.2921},\n",
       "                                {'token': ' perform', 'weight': 0.4293},\n",
       "                                {'token': ' well', 'weight': 0.2772},\n",
       "                                {'token': '.', 'weight': 4.4723}],\n",
       "               'score': 0.12305451184511185,\n",
       "               'text': 'Their solution did not perform well.'},\n",
       "   'tags': ['find_highlights', 'online', 'development'],\n",
       "   'trace_id': '7fcf8271-1738-4025-8305-d5a1e5100aea'}))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    find_highlights(\n",
    "        \"Our solution has outperformed the state-of-the-art.\"\n",
    "    ), find_highlights(\"Their solution did not perform well.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, the service is built as a docker image, pushed to our image registry and subsequent rolling update is performed in the production cluster.\n",
    "To check out the Dockerimage, go to [the additional files page](/examples/scibert/additional-files)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('.env': venv)",
   "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.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "02dd6d3afbfa9fbbe1037d64ad9014965528a1ccad21929d6e72f466389a68ad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}