300 lines
11 KiB
Text
300 lines
11 KiB
Text
{
|
|
"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<?, ?it/s]"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"score=0.4977932274341583 text='Our solution has outperformed the state-of-the-art.' explanation=[Attention(weight=0.1332, token='our'), Attention(weight=0.133, token='[UNK]'), Attention(weight=0.1332, token='[UNK]'), Attention(weight=0.1329, token='[UNK]'), Attention(weight=0.134, token='the'), Attention(weight=0.1342, token='[UNK]'), Attention(weight=0.1333, token='[UNK]'), Attention(weight=0.1332, token='[UNK]'), Attention(weight=0.1327, token='[UNK]'), Attention(weight=0.1316, token='the'), Attention(weight=0.1339, token='[UNK]'), Attention(weight=0.1325, token='[UNK]'), Attention(weight=0.1347, token='.')]\n",
|
|
"score=0.49779173731803894 text='Their solution did not perform well.' explanation=[Attention(weight=0.2217, token='[UNK]'), Attention(weight=0.2211, token='[UNK]'), Attention(weight=0.2215, token='[UNK]'), Attention(weight=0.221, token='[UNK]'), Attention(weight=0.2239, token='[UNK]'), Attention(weight=0.223, token='[UNK]'), Attention(weight=0.2243, token='.')]\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"if __name__ == \"__main__\":\n",
|
|
" from fastapi.testclient import TestClient\n",
|
|
"\n",
|
|
" first = find_highlights(\"Our solution has outperformed the state-of-the-art.\")\n",
|
|
" second = find_highlights(\"Their solution did not perform well.\")\n",
|
|
" assert 0 <= first.output.score <= 1\n",
|
|
" assert first.output.explanation\n",
|
|
" print(first.output)\n",
|
|
" print(second.output)\n",
|
|
" with TestClient(find_highlights.app) as client:\n",
|
|
" response = client.post(\n",
|
|
" \"/predict\", json={\"sentence\": \"Our method improves accuracy.\"}\n",
|
|
" )\n",
|
|
" response.raise_for_status()\n",
|
|
" assert 0 <= response.json()[\"output\"][\"score\"] <= 1"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cell-12",
|
|
"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](https://great-ai.schmelczer.dev/examples/scibert/additional-files)."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv (3.13.5)",
|
|
"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.13.5"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|