{ "cells": [ { "cell_type": "markdown", "id": "cell-0", "metadata": {}, "source": [ "# Create an inference function\n", "\n", "Run `data.ipynb` and `train.ipynb` first in the same local/full mode.\n", "This notebook loads their local model and tokenizer; no S3 credentials,\n", "MongoDB server, or model downloads are needed for the default demo.\n", "The tiny demo model's scores have no scientific interpretation." ] }, { "cell_type": "markdown", "id": "85c2ee1d", "metadata": {}, "source": [ "## Local setup\n", "\n", "Use Python 3.12 or 3.13 and install this checkout with `uv pip install --editable '.[notebooks]'`.\n", "See **Run notebooks locally** in the repository README for the kernel, execution order,\n", "and full-data options. Run related notebooks from the same working directory\n", "to share the default `.cache/` model store and `tracing_database.json` database.\n" ] }, { "cell_type": "markdown", "id": "cell-4", "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": 1, "id": "cell-5", "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:48:09.038647Z", "iopub.status.busy": "2026-10-28T22:48:09.038493Z", "iopub.status.idle": "2026-10-28T22:48:09.083813Z", "shell.execute_reply": "2026-10-28T22:48:09.083283Z" } }, "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", "id": "cell-6", "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": 2, "id": "cell-7", "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:48:09.084713Z", "iopub.status.busy": "2026-10-28T22:48:09.084644Z", "iopub.status.idle": "2026-10-28T22:48:11.256390Z", "shell.execute_reply": "2026-10-28T22:48:11.256030Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;226mEnvironment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;226mCannot find credentials files, defaulting to using ParallelTinyDbDriver\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;226mThe selected tracing database (ParallelTinyDbDriver) is not recommended for production\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;226mCannot find credentials files, defaulting to using LargeFileLocal\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39mGreatAI (v0.1.11): configured ✅\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39m 🔩 tracing_database: ParallelTinyDbDriver\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39m 🔩 large_file_implementation: LargeFileLocal\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39m 🔩 is_production: False\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39m 🔩 should_log_exception_stack: True\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39m 🔩 prediction_cache_size: 512\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39m 🔩 dashboard_table_size: 50\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;226mYou still need to check whether you follow all best practices before trusting your deployment.\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;226m> Find out more at https://se-ml.github.io/practices\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39mFetching cached versions of scibert-highlights\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39mLatest version of scibert-highlights is 0 (from versions: 0)\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39mFile scibert-highlights-0 found in cache\u001b[0m\n" ] } ], "source": [ "from functools import lru_cache\n", "from pathlib import Path\n", "from great_ai import use_model\n", "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n", "\n", "\n", "@lru_cache(maxsize=1)\n", "def load_tokenizer_and_model(model_path):\n", " tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)\n", " model = AutoModelForSequenceClassification.from_pretrained(\n", " model_path, local_files_only=True, attn_implementation=\"eager\"\n", " )\n", " model.eval()\n", " return tokenizer, model\n", "\n", "\n", "@use_model(\"scibert-highlights\", version=\"latest\", model_kwarg_name=\"model_path\")\n", "def get_tokenizer_and_model(model_path: Path):\n", " return load_tokenizer_and_model(model_path)" ] }, { "cell_type": "markdown", "id": "cell-8", "metadata": {}, "source": [ "Finally, implement the inference function." ] }, { "cell_type": "code", "execution_count": 3, "id": "cell-9", "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:48:11.257418Z", "iopub.status.busy": "2026-10-28T22:48:11.257269Z", "iopub.status.idle": "2026-10-28T22:48:11.289789Z", "shell.execute_reply": "2026-10-28T22:48:11.289400Z" } }, "outputs": [], "source": [ "from great_ai import GreatAI\n", "from great_ai.utilities import clean\n", "import torch\n", "\n", "\n", "@GreatAI.create\n", "def find_highlights(sentence: str) -> EvaluatedSentence:\n", " \"\"\"Score a sentence and expose the final layer's CLS attention per token.\n", "\n", " Attention weights are a model diagnostic, not a causal explanation.\n", " \"\"\"\n", " tokenizer, model = get_tokenizer_and_model()\n", " sentence = clean(sentence, convert_to_ascii=True, remove_brackets=True)\n", " tensors = tokenizer(sentence, return_tensors=\"pt\", truncation=True, max_length=512)\n", " with torch.inference_mode():\n", " result = model(**tensors, output_attentions=True)\n", " score = torch.softmax(result.logits, dim=1)[0, 1].item()\n", " weights = result.attentions[-1][0, :, 0, :].sum(dim=0).cpu().tolist()\n", " tokens = tokenizer.convert_ids_to_tokens(tensors[\"input_ids\"][0].tolist())\n", " explanation = [\n", " Attention(token=token, weight=round(weight, 4))\n", " for token, weight in zip(tokens[1:-1], weights[1:-1])\n", " ]\n", " return EvaluatedSentence(score=score, text=sentence, explanation=explanation)" ] }, { "cell_type": "markdown", "id": "cell-10", "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": 4, "id": "cell-11", "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:48:11.290551Z", "iopub.status.busy": "2026-10-28T22:48:11.290495Z", "iopub.status.idle": "2026-10-28T22:48:11.803477Z", "shell.execute_reply": "2026-10-28T22:48:11.803126Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "W1028 22:48:11.714000 21400 .venv/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/redirects.py:35] NOTE: Redirects are currently not supported in MacOs.\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "494840f6abbc4b7aaad2ce762d1e32e8", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Loading weights: 0%| | 0/25 [00:00