{ "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. Generated artifacts stay under `.cache/notebooks/`.\n" ] }, { "cell_type": "code", "execution_count": 9, "id": "3775e11d", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;196mConfiguration has been already initialised, overwriting.\n", "Make sure to call `configure()` before importing your application code.\u001b[0m\n", "\u001b[38;5;226mEnvironment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n", "\u001b[38;5;226mThe selected tracing database (ParallelTinyDbDriver) is not recommended for production\u001b[0m\n", "\u001b[38;5;39mGreatAI (v0.1.11): configured ✅\u001b[0m\n", "\u001b[38;5;39m 🔩 tracing_database: ParallelTinyDbDriver\u001b[0m\n", "\u001b[38;5;39m 🔩 large_file_implementation: LargeFileLocal\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: 512\u001b[0m\n", "\u001b[38;5;39m 🔩 dashboard_table_size: 50\u001b[0m\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Notebook artifacts: /Users/schmelczer/Projects/great-ai/.cache/notebooks/scibert\n" ] } ], "source": [ "import os\n", "import sys\n", "from pathlib import Path\n", "\n", "# Works when opened from the repository root or any notebook directory.\n", "_start = Path(globals().get(\"__file__\", Path.cwd())).resolve()\n", "ROOT = next(\n", " p\n", " for p in [_start, *_start.parents]\n", " if (p / \"pyproject.toml\").is_file() and (p / \"great_ai\").is_dir()\n", ")\n", "sys.path.insert(0, str(ROOT / \"docs\"))\n", "from notebook_support import setup_notebook, mag_lines\n", "\n", "FULL_RUN = os.environ.get(\"GREAT_AI_NOTEBOOK_FULL\") == \"1\"\n", "WORK_DIR = setup_notebook(ROOT, \"scibert\")\n", "DATA_DIR = ROOT / \"docs/tutorial/data\"" ] }, { "cell_type": "code", "execution_count": 10, "id": "cell-3", "metadata": {}, "outputs": [], "source": [ "# The setup cell selects local model storage and a local TinyDB tracing database." ] }, { "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": 11, "id": "cell-5", "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", "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": 12, "id": "cell-7", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39mFetching cached versions of scibert-highlights\u001b[0m\n", "\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 functools import lru_cache\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": 13, "id": "cell-9", "metadata": {}, "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": 14, "id": "cell-11", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "W1028 12:49:11.667000 52910 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": "6ef4b00220304248a94035042cb5fa15", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Loading weights: 0%| | 0/25 [00:00