005_module_ml.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a19ce1-bb8a-40c1-8ef3-ff78f60b71d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp ml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885c918b-5721-4a0d-947e-62451a6d4c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec96dcb-56ee-412f-bc63-253c5b1a19b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from nbdev import nbdev_export\n",
    "from fastcore.test import test_eq, test_fail"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f00408-3ef3-475c-8ab4-eea43088ecae",
   "metadata": {},
   "source": [
    "# ml\n",
    "> Reference API related to ml applications"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e79fb1c-ff24-4887-842d-6230f7c599f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/tmartins/.local/share/virtualenvs/learntorank-qPFjAqDO/lib/python3.9/site-packages/tqdm/auto.py:22: 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",
      "2022-11-04 09:54:26.714852: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "#|export\n",
    "import sys\n",
    "from os import PathLike\n",
    "from typing import List, Optional, Union, Mapping, Dict, IO\n",
    "from pathlib import Path\n",
    "from urllib.parse import urlencode\n",
    "from fastcore.utils import patch, patch_to\n",
    "\n",
    "from vespa.package import (\n",
    "    OnnxModel,\n",
    "    QueryTypeField,\n",
    "    Field,\n",
    "    Function,\n",
    "    RankProfile,\n",
    "    ApplicationPackage\n",
    ")\n",
    "from torch import tensor\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification,\n",
    "    BertForSequenceClassification,\n",
    "    BertTokenizerFast,\n",
    "    pipeline,\n",
    ")\n",
    "from transformers.convert_graph_to_onnx import convert_pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c079c56-18ec-4221-9245-3384e68cd431",
   "metadata": {},
   "source": [
    "## ML tasks involving text inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d16c13d-811a-4c47-b878-8c183dcd3a0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Task(object):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_id: str,  # Id used to identify the model on Vespa applications.\n",
    "    ):\n",
    "        \"Base class for ML Tasks.\"\n",
    "        self.model_id = model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bee5fc41-b5bc-4c02-beef-509e5026437a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class TextTask(Task):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_id: str,  # Id used to identify the model on Vespa applications.\n",
    "        model: str,  # Id of the model as used by the model hub.\n",
    "        tokenizer: Optional[str] = None,  # Id of the tokenizer as used by the model hub.\n",
    "        output_file: IO = sys.stdout,  # Output file to write output messages.\n",
    "    ):\n",
    "        \"Base class for Tasks involving text inputs.\"\n",
    "        super().__init__(model_id=model_id)\n",
    "        self.model = model\n",
    "        self.tokenizer = tokenizer\n",
    "        if not self.tokenizer:\n",
    "            self.tokenizer = model\n",
    "        self.output = output_file\n",
    "        self._tokenizer = None\n",
    "        self._model = None\n",
    "\n",
    "    @staticmethod\n",
    "    def parse_vespa_prediction(prediction) -> List:\n",
    "        return prediction[\"values\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac224e3d-dcc8-4233-818c-8b50ad922434",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _load_tokenizer(self: TextTask):\n",
    "    if not self._tokenizer:\n",
    "        print(\"Downloading tokenizer.\", file=self.output)\n",
    "        self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bd87ad7-4d2d-4eca-9c73-8e993e1838ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _create_pipeline(self: TextTask):\n",
    "    self._load_tokenizer()\n",
    "    if not self._model:\n",
    "        print(\"Downloading model.\", file=self.output)\n",
    "        self._model = AutoModelForSequenceClassification.from_pretrained(self.model)\n",
    "        print(\"Model loaded.\", file=self.output)\n",
    "\n",
    "    _pipeline = pipeline(\n",
    "        task=\"text-classification\",\n",
    "        model=self._model,\n",
    "        tokenizer=self._tokenizer,\n",
    "        top_k=2,\n",
    "        function_to_apply=\"None\",\n",
    "    )\n",
    "    return _pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3135d7a5-52bb-402e-a53b-b3414bc7d23a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def export_to_onnx(\n",
    "    self: TextTask, \n",
    "    output_path: str  # Relative output path for the onnx model, should end in '.onnx'\n",
    ") -> None:\n",
    "    \"Export a model to ONNX\"\n",
    "    pipeline = self._create_pipeline()\n",
    "    convert_pytorch(\n",
    "        pipeline, opset=11, output=Path(output_path), use_external_format=False\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5167b79a-b4e7-4273-801f-28a291c75ff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def predict(\n",
    "    self: TextTask, \n",
    "    text: str  # text input for the task.\n",
    ") -> List:  # Predictions.\n",
    "    \"Predict using a local instance of the model\"\n",
    "    pipeline = self._create_pipeline()\n",
    "    predictions = pipeline(text)\n",
    "    return [\n",
    "        [x[\"score\"] for x in predictions[0] if x[\"label\"] == \"LABEL_0\"][0],\n",
    "        [x[\"score\"] for x in predictions[0] if x[\"label\"] == \"LABEL_1\"][0],\n",
    "    ]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f906dabe-9df4-45d0-b365-8db3bc479017",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_url_encoded_tokens(self: TextTask, x):\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd126d36-1fe8-4baa-971d-c2276cb98da0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class SequenceClassification(TextTask):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_id: str,  # Id used to identify the model on Vespa applications.\n",
    "        model: str,  # Id of the model as used by the model hub. Alternatively, it can also be the path to the folder containing the model files, as long as the model config is also there.\n",
    "        tokenizer: Optional[str] = None,  # Id of the tokenizer as used by the model hub. Alternatively, it can also be the path to the folder containing the tokenizer files, as long as the model config is also there.\n",
    "        output_file: IO = sys.stdout,  # Output file to write output messages.\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Sequence Classification task.\n",
    "\n",
    "        It takes a text input and returns an array of floats depending on which\n",
    "        model is used to solve the task.\n",
    "        \"\"\"\n",
    "        super().__init__(\n",
    "            model_id=model_id, model=model, tokenizer=tokenizer, output_file=output_file\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac92fdf-7506-4d4c-a19c-e9300d0dfc76",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_url_encoded_tokens(\n",
    "    self: SequenceClassification, \n",
    "    x\n",
    "):\n",
    "    self._load_tokenizer()\n",
    "    tokens = self._tokenizer(x)\n",
    "    encoded_tokens = urlencode(\n",
    "        {\n",
    "            key: \"{\"\n",
    "            + \",\".join(\n",
    "                [\n",
    "                    \"{{d0: 0, d1: {}}}: {}\".format(idx, x)\n",
    "                    for idx, x in enumerate(value)\n",
    "                ]\n",
    "            )\n",
    "            + \"}\"\n",
    "            for key, value in tokens.items()\n",
    "        }\n",
    "    )\n",
    "    return encoded_tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "731a92a2-472a-4a26-bc0b-322c271c5c64",
   "metadata": {},
   "source": [
    "## Model config for Vespa applications"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909a62df-5ae9-40e9-9eaf-f753c9e7b465",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ModelConfig(object):\n",
    "    def __init__(\n",
    "        self, \n",
    "        model_id  # Unique model id to represent the model within a Vespa application.\n",
    "    ) -> None:\n",
    "        \"Base model configuration for Vespa applications.\"\n",
    "        \n",
    "        self.model_id = model_id\n",
    "\n",
    "    def onnx_model(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def query_profile_type_fields(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def document_fields(self, document_field_indexing):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def rank_profile(self, include_model_summary_features, **kwargs):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7afad2e0-f553-4784-9055-e57f6cf0392b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class BertModelConfig(ModelConfig):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_id: str,  # Unique model id to represent the model within a Vespa application.\n",
    "        query_input_size: int,  # The size of the input vector dedicated to the query text.\n",
    "        doc_input_size: int,  # The size of the input vector dedicated to the document text.\n",
    "        tokenizer: Union[str, PathLike],  # The name or a path to a saved BERT model tokenizer from the transformers library.\n",
    "        model: Optional[Union[str, PathLike]] = None,  # The name or a path to a saved model that is compatible with the `tokenizer`. The model is optional at construction since you might want to train it first. You must add a model via :func:`add_model` before deploying a Vespa application that uses this class.\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        BERT model configuration for Vespa applications.\n",
    "\n",
    "\n",
    "        >>> bert_config = BertModelConfig(\n",
    "        ...     model_id=\"pretrained_bert_tiny\",\n",
    "        ...     query_input_size=32,\n",
    "        ...     doc_input_size=96,\n",
    "        ...     tokenizer=\"google/bert_uncased_L-2_H-128_A-2\",\n",
    "        ...     model=\"google/bert_uncased_L-2_H-128_A-2\",\n",
    "        ... )  # doctest: +SKIP\n",
    "        BertModelConfig('pretrained_bert_tiny', 32, 96, 'google/bert_uncased_L-2_H-128_A-2', 'google/bert_uncased_L-2_H-128_A-2')\n",
    "        \"\"\"\n",
    "        super().__init__(model_id=model_id)\n",
    "\n",
    "        self.query_input_size = query_input_size\n",
    "        self.doc_input_size = doc_input_size\n",
    "        self.input_size = self.query_input_size + self.doc_input_size\n",
    "\n",
    "        self.tokenizer = tokenizer\n",
    "        self._tokenizer = BertTokenizerFast.from_pretrained(tokenizer)\n",
    "        self._validate_tokenizer()\n",
    "\n",
    "        self.query_token_ids_name = model_id + \"_query_token_ids\"\n",
    "        self.actual_query_input_size = (\n",
    "            query_input_size - 2\n",
    "        )  # one character saved for CLS and one for SEP\n",
    "\n",
    "        self.doc_token_ids_name = model_id + \"_doc_token_ids\"\n",
    "        self.actual_doc_input_size = doc_input_size - 1  # one character saved for SEP\n",
    "\n",
    "        self.model = model\n",
    "        self._model = None\n",
    "        if model:\n",
    "            self.add_model(model=model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ce37a93-cd06-4639-97d6-bae80980a588",
   "metadata": {},
   "outputs": [],
   "source": [
    "#bert_config = BertModelConfig(\n",
    "#    model_id=\"pretrained_bert_tiny\",\n",
    "#    query_input_size=32,\n",
    "#    doc_input_size=96,\n",
    "#    tokenizer=\"google/bert_uncased_L-2_H-128_A-2\"\n",
    "#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af302389-3859-481b-8599-37b0be2fea30",
   "metadata": {},
   "outputs": [],
   "source": [
    "#bert_config = BertModelConfig(\n",
    "#    model_id=\"pretrained_bert_tiny\",\n",
    "#    query_input_size=32,\n",
    "#    doc_input_size=96,\n",
    "#    tokenizer=\"google/bert_uncased_L-2_H-128_A-2\",\n",
    "#    model=\"google/bert_uncased_L-2_H-128_A-2\",\n",
    "#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b74c2503-c31b-4a85-98dc-3a90e0bd91aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def predict(\n",
    "    self: BertModelConfig, \n",
    "    queries,  # A List of query texts. \n",
    "    docs  # A List of document texts.\n",
    ") -> List:  # Logits\n",
    "    \"Predict (forward pass) given queries and docs texts\"\n",
    "\n",
    "    if not self._model:\n",
    "        raise ValueError(\"A model needs to be added.\")\n",
    "    model_output = self._model(\n",
    "        **self.create_encodings(queries=queries, docs=docs, return_tensors=True),\n",
    "        return_dict=True\n",
    "    )\n",
    "    return model_output.logits.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0743eacf-71bd-4c0a-8379-31587bda3eed",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _validate_tokenizer(self:BertModelConfig) -> None:\n",
    "    dummy_inputs = self._generate_dummy_inputs()\n",
    "\n",
    "    assert (\n",
    "        dummy_inputs[\"input_ids\"].shape[1]\n",
    "        == dummy_inputs[\"token_type_ids\"].shape[1]\n",
    "        and dummy_inputs[\"token_type_ids\"].shape[1]\n",
    "        == dummy_inputs[\"attention_mask\"].shape[1]\n",
    "        and dummy_inputs[\"attention_mask\"].shape[1] == self.input_size\n",
    "    ), \"tokenizer generates wrong input size\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d478534d-f44c-43cc-babb-49a49732d042",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _validate_model(self: BertModelConfig, model: BertForSequenceClassification) -> None:\n",
    "    if not isinstance(model, BertForSequenceClassification):\n",
    "        raise ValueError(\"We only support BertForSequenceClassification for now.\")\n",
    "    model_output = model(**self._generate_dummy_inputs(), return_dict=True)\n",
    "    if len(model_output.logits.shape) != 2:\n",
    "        ValueError(\"Model output expected to be logits vector of size 2\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ea9db11-723e-4186-ad54-ff732da1d588",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def add_model(\n",
    "    self: BertModelConfig, \n",
    "    model: Union[str, PathLike]  # The name or a path to a saved model that is compatible with the `tokenizer`.\n",
    ") -> None:\n",
    "    \"Add a BERT model\"\n",
    "    \n",
    "    _model = BertForSequenceClassification.from_pretrained(model)\n",
    "    self._validate_model(model=_model)\n",
    "    self.model = model\n",
    "    self._model = _model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953a64d6-ce3d-4585-9333-201f1ec559a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _query_input_ids(self: BertModelConfig, queries: List[str]):\n",
    "    queries_encodings = self._tokenizer(\n",
    "        queries,\n",
    "        truncation=True,\n",
    "        max_length=self.query_input_size - 2,\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "    return queries_encodings[\"input_ids\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83480ded-348b-40cd-8311-0b10a8c76a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _doc_input_ids(self: BertModelConfig, docs: List[str]):\n",
    "    docs_encodings = self._tokenizer(\n",
    "        docs,\n",
    "        truncation=True,\n",
    "        max_length=self.doc_input_size - 1,\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "    return docs_encodings[\"input_ids\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1684fc92-6817-4a74-8d40-52acdcfae1ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def doc_fields(\n",
    "    self: BertModelConfig, \n",
    "    text: str  # The text related to the document to be used as input to the bert model\n",
    ") -> Dict:  # Dict with key and values as expected by Vespa.\n",
    "    \"Generate document fields related to the model that needs to be fed to Vespa.\"\n",
    "    \n",
    "    input_ids = self._doc_input_ids([text])[0]\n",
    "    if len(input_ids) < self.actual_doc_input_size:\n",
    "        input_ids = input_ids + [0] * (self.actual_doc_input_size - len(input_ids))\n",
    "    return {self.doc_token_ids_name: {\"values\": input_ids}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ee6104a-360c-4f16-9a23-2c1f68c0733c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def query_tensor_mapping(\n",
    "    self: BertModelConfig, \n",
    "    text: str # Query text to be used as input to the BERT model.\n",
    ") -> List[float]:  # Input ids expected by Vespa.\n",
    "    \"Maps query text to a tensor expected by Vespa at run time.\"\n",
    "    \n",
    "    input_ids = self._query_input_ids([text])[0]\n",
    "    if len(input_ids) < self.actual_query_input_size:\n",
    "        input_ids = input_ids + [0] * (\n",
    "            self.actual_query_input_size - len(input_ids)\n",
    "        )\n",
    "    return input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd0892dc-8b67-4526-a344-57575310372d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def create_encodings(\n",
    "    self: BertModelConfig, \n",
    "    queries: List[str], # Query texts.\n",
    "    docs: List[str],  # Document texts. \n",
    "    return_tensors=False  # Return tensors\n",
    ") -> Dict:  # Dict containing `input_ids`, `token_type_ids` and `attention_mask` encodings.\n",
    "    \"\"\"\n",
    "    Create BERT model encodings.\n",
    "\n",
    "    Create BERT encodings following the same pattern used during Vespa serving. Useful to generate training data\n",
    "    and ensuring training and serving compatibility.\n",
    "    \"\"\"\n",
    "    query_input_ids = self._query_input_ids(queries=queries)\n",
    "    doc_input_ids = self._doc_input_ids(docs=docs)\n",
    "\n",
    "    TOKEN_NONE = 0\n",
    "    TOKEN_CLS = 101\n",
    "    TOKEN_SEP = 102\n",
    "\n",
    "    input_ids = []\n",
    "    token_type_ids = []\n",
    "    attention_mask = []\n",
    "    for query_input_id, doc_input_id in zip(query_input_ids, doc_input_ids):\n",
    "        # create input id\n",
    "        input_id = (\n",
    "            [TOKEN_CLS] + query_input_id + [TOKEN_SEP] + doc_input_id + [TOKEN_SEP]\n",
    "        )\n",
    "        number_tokens = len(input_id)\n",
    "        padding_length = max(self.input_size - number_tokens, 0)\n",
    "        input_id = input_id + [TOKEN_NONE] * padding_length\n",
    "        input_ids.append(input_id)\n",
    "        # create token id\n",
    "        token_type_id = (\n",
    "            [0] * len([TOKEN_CLS] + query_input_id + [TOKEN_SEP])\n",
    "            + [1] * len(doc_input_id + [TOKEN_SEP])\n",
    "            + [TOKEN_NONE] * padding_length\n",
    "        )\n",
    "        token_type_ids.append(token_type_id)\n",
    "        # create attention_mask\n",
    "        attention_mask.append([1] * number_tokens + [TOKEN_NONE] * padding_length)\n",
    "\n",
    "    if return_tensors:\n",
    "        input_ids = tensor(input_ids)\n",
    "        token_type_ids = tensor(token_type_ids)\n",
    "        attention_mask = tensor(attention_mask)\n",
    "    encodings = {\n",
    "        \"input_ids\": input_ids,\n",
    "        \"token_type_ids\": token_type_ids,\n",
    "        \"attention_mask\": attention_mask,\n",
    "    }\n",
    "    return encodings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c11b6158-4d7f-41ae-b0d0-3933b5ade79d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def _generate_dummy_inputs(self:BertModelConfig):\n",
    "    dummy_input = self.create_encodings(\n",
    "        queries=[\"dummy query 1\"], docs=[\"dummy document 1\"], return_tensors=True\n",
    "    )\n",
    "    return dummy_input\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a548b085-374b-485e-a7fa-d92386f9ce8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def export_to_onnx(\n",
    "    self: BertModelConfig, \n",
    "    output_path: str  # Relative output path for the onnx model, should end in '.onnx'\n",
    ") -> None:\n",
    "    \"Export a model to ONNX\"\n",
    "\n",
    "    if self._model:\n",
    "        _pipeline = pipeline(\n",
    "            task=\"text-classification\",\n",
    "            model=self._model,\n",
    "            tokenizer=self._tokenizer,\n",
    "            return_all_scores=True,\n",
    "            function_to_apply=\"None\",\n",
    "        )\n",
    "        convert_pytorch(\n",
    "            _pipeline, opset=11, output=Path(output_path), use_external_format=False\n",
    "        )\n",
    "    else:\n",
    "        raise ValueError(\"No BERT model found to be exported.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa65a8b-0f2e-430c-97f1-655e31ca0a43",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def onnx_model(self: BertModelConfig):\n",
    "    model_file_path = self.model_id + \".onnx\"\n",
    "    self.export_to_onnx(output_path=model_file_path)\n",
    "\n",
    "    return OnnxModel(\n",
    "        model_name=self.model_id,\n",
    "        model_file_path=model_file_path,\n",
    "        inputs={\n",
    "            \"input_ids\": \"input_ids\",\n",
    "            \"token_type_ids\": \"token_type_ids\",\n",
    "            \"attention_mask\": \"attention_mask\",\n",
    "        },\n",
    "        outputs={\"output_0\": \"logits\"},\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b34fb83-32ff-4631-9a01-d596c8654156",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def query_profile_type_fields(self: BertModelConfig):\n",
    "    return [\n",
    "        QueryTypeField(\n",
    "            name=\"ranking.features.query({})\".format(self.query_token_ids_name),\n",
    "            type=\"tensor<float>(d0[{}])\".format(int(self.actual_query_input_size)),\n",
    "        )\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c8bab82-0eda-4bd9-adbc-7692c0ff9804",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def document_fields(self: BertModelConfig, document_field_indexing):\n",
    "    if not document_field_indexing:\n",
    "        document_field_indexing = [\"attribute\", \"summary\"]\n",
    "\n",
    "    return [\n",
    "        Field(\n",
    "            name=self.doc_token_ids_name,\n",
    "            type=\"tensor<float>(d0[{}])\".format(int(self.actual_doc_input_size)),\n",
    "            indexing=document_field_indexing,\n",
    "        ),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25dc8372-ee9d-4555-86d2-06dfeda2a78a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def rank_profile(self: BertModelConfig, include_model_summary_features, **kwargs):\n",
    "    constants = {\"TOKEN_NONE\": 0, \"TOKEN_CLS\": 101, \"TOKEN_SEP\": 102}\n",
    "    # ToDo: Add a unit test for constants\n",
    "    if \"constants\" in kwargs:\n",
    "        constants.update(kwargs.pop(\"constants\"))\n",
    "\n",
    "    functions = [\n",
    "        Function(\n",
    "            name=\"question_length\",\n",
    "            expression=\"sum(map(query({}), f(a)(a > 0)))\".format(\n",
    "                self.query_token_ids_name\n",
    "            ),\n",
    "        ),\n",
    "        Function(\n",
    "            name=\"doc_length\",\n",
    "            expression=\"sum(map(attribute({}), f(a)(a > 0)))\".format(\n",
    "                self.doc_token_ids_name\n",
    "            ),\n",
    "        ),\n",
    "        Function(\n",
    "            name=\"input_ids\",\n",
    "            expression=\"tokenInputIds({}, query({}), attribute({}))\".format(\n",
    "                self.input_size,\n",
    "                self.query_token_ids_name,\n",
    "                self.doc_token_ids_name,\n",
    "            ),\n",
    "        ),\n",
    "        Function(\n",
    "            name=\"attention_mask\",\n",
    "            expression=\"tokenAttentionMask({}, query({}), attribute({}))\".format(\n",
    "                self.input_size,\n",
    "                self.query_token_ids_name,\n",
    "                self.doc_token_ids_name,\n",
    "            ),\n",
    "        ),\n",
    "        Function(\n",
    "            name=\"token_type_ids\",\n",
    "            expression=\"tokenTypeIds({}, query({}), attribute({}))\".format(\n",
    "                self.input_size,\n",
    "                self.query_token_ids_name,\n",
    "                self.doc_token_ids_name,\n",
    "            ),\n",
    "        ),\n",
    "        Function(\n",
    "            name=\"logit0\",\n",
    "            expression=\"onnx(\" + self.model_id + \").logits{d0:0,d1:0}\",\n",
    "        ),\n",
    "        Function(\n",
    "            name=\"logit1\",\n",
    "            expression=\"onnx(\" + self.model_id + \").logits{d0:0,d1:1}\",\n",
    "        ),\n",
    "    ]\n",
    "    if \"functions\" in kwargs:\n",
    "        functions.extend(kwargs.pop(\"functions\"))\n",
    "\n",
    "    summary_features = []\n",
    "    if include_model_summary_features:\n",
    "        summary_features.extend(\n",
    "            [\n",
    "                \"logit0\",\n",
    "                \"logit1\",\n",
    "                \"input_ids\",\n",
    "                \"attention_mask\",\n",
    "                \"token_type_ids\",\n",
    "            ]\n",
    "        )\n",
    "    if \"summary_features\" in kwargs:\n",
    "        summary_features.extend(kwargs.pop(\"summary_features\"))\n",
    "\n",
    "    return RankProfile(\n",
    "        name=self.model_id,\n",
    "        constants=constants,\n",
    "        functions=functions,\n",
    "        summary_features=summary_features,\n",
    "        **kwargs\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4929ea33-6827-474d-a3f6-a331cf6bd2ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def __eq__(self: BertModelConfig, other):\n",
    "    if not isinstance(other, self.__class__):\n",
    "        return False\n",
    "    return (\n",
    "        self.model_id == other.model_id\n",
    "        and self.query_input_size == other.query_input_size\n",
    "        and self.doc_input_size == other.doc_input_size\n",
    "        and self.tokenizer == other.tokenizer\n",
    "        and self.model == other.model\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f77825d-d2f6-4fc8-a12e-4fa77788013e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def __repr__(self: BertModelConfig):\n",
    "    return \"{0}({1}, {2}, {3}, {4}, {5})\".format(\n",
    "        self.__class__.__name__,\n",
    "        repr(self.model_id),\n",
    "        repr(self.query_input_size),\n",
    "        repr(self.doc_input_size),\n",
    "        repr(self.tokenizer),\n",
    "        repr(self.model),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82be2f6-ffb0-4d15-b0b9-e4bb1b395567",
   "metadata": {},
   "source": [
    "## Model Server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da2d02ea-eef2-4641-8326-dd2b96186e1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ModelServer(ApplicationPackage):\n",
    "    \"\"\"\n",
    "    Create a Vespa stateless model evaluation server.\n",
    "\n",
    "    A Vespa stateless model evaluation server is a simplified Vespa application without content clusters.\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        name: str,  # Application name.\n",
    "        tasks: Optional[List[Task]] = None,  # List of tasks to be served.\n",
    "    ):\n",
    "        super().__init__(\n",
    "            name=name,\n",
    "            schema=None,\n",
    "            query_profile=None,\n",
    "            query_profile_type=None,\n",
    "            stateless_model_evaluation=True,\n",
    "            create_schema_by_default=False,\n",
    "            create_query_profile_by_default=False,\n",
    "        )\n",
    "        self.models = {} if not tasks else {model.model_id: model for model in tasks}\n",
    "\n",
    "    @staticmethod\n",
    "    def from_dict(mapping: Mapping) -> \"ModelServer\":\n",
    "        return ModelServer(name=mapping[\"name\"])\n",
    "\n",
    "    @property\n",
    "    def to_dict(self) -> Mapping:\n",
    "        return {\"name\": self.name}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c358e273-d4bc-44ad-a0ed-d2f141d45f6a",
   "metadata": {},
   "source": [
    "## Add ranking model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da0d3a59-bcea-4f3d-b900-7b36bedb64ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def add_ranking_model(\n",
    "    app_package: ApplicationPackage,  # Application package to include ranking model\n",
    "    model_config: ModelConfig,  # Model config instance specifying the model to be used on the RankProfile.\n",
    "    schema=None,  # Name of the schema to add model ranking to.\n",
    "    include_model_summary_features=False,  # True to include model specific summary features, such as inputs and outputs that are useful for debugging. Default to False as this requires an extra model evaluation when fetching summary features.\n",
    "    document_field_indexing=None,  # List of indexing attributes for the document fields required by the ranking model.\n",
    "    **kwargs  # Further arguments to be passed to `RankProfile`.\n",
    ") -> None:\n",
    "    \"Add ranking profile based on a specific model config.\"\n",
    "\n",
    "    model_id = model_config.model_id\n",
    "    #\n",
    "    # Validate and persist config\n",
    "    #\n",
    "    if model_id in app_package.model_ids:\n",
    "        raise ValueError(\"model_id must be unique: {}\".format(model_id))\n",
    "    app_package.model_ids.append(model_id)\n",
    "    app_package.model_configs[model_id] = model_config\n",
    "    #\n",
    "    # Export ONNX model\n",
    "    #\n",
    "    app_package.get_schema(schema).add_model(model_config.onnx_model())\n",
    "    #\n",
    "    # Add query profile type fields\n",
    "    #\n",
    "    app_package.query_profile_type.add_fields(*model_config.query_profile_type_fields())\n",
    "    #\n",
    "    # Add field for doc token ids\n",
    "    #\n",
    "    app_package.get_schema(schema).add_fields(\n",
    "        *model_config.document_fields(\n",
    "            document_field_indexing=document_field_indexing\n",
    "        )\n",
    "    )\n",
    "    #\n",
    "    # Add rank profiles\n",
    "    #\n",
    "    app_package.get_schema(schema).add_rank_profile(\n",
    "        model_config.rank_profile(\n",
    "            include_model_summary_features=include_model_summary_features, **kwargs\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d77781-8f0d-46d0-92f6-22ef8b25719d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "learntorank",
   "language": "python",
   "name": "learntorank"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}