diff --git a/custom-models/README.md b/custom-models/README.md index e20465fd..588566e4 100644 --- a/custom-models/README.md +++ b/custom-models/README.md @@ -24,6 +24,7 @@ Before creating a custom model job (fine-tuning or distillation) in Amazon Bedro - [Fine-Tuning](./bedrock-fine-tuning) - Exmaple showing how different LLMs can be fine tuned onto Amazon Bedrock. - [Model Distillation](./model_distillation) - Exmaple showing how different LLMs can be distilled onto Amazon Bedrock. - [import_models](./import_models) - Exmaple showing how different Opensource LLMs can be imported onto Amazon Bedrock. +- [Structured Output with CMI](./structured_output_with_cmi.ipynb) - Example showing how to constrain imported model outputs to a JSON schema using `response_format`. ## Contributing diff --git a/custom-models/structured_output_with_cmi.ipynb b/custom-models/structured_output_with_cmi.ipynb new file mode 100644 index 00000000..e76035f7 --- /dev/null +++ b/custom-models/structured_output_with_cmi.ipynb @@ -0,0 +1,479 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Structured Output with Amazon Bedrock Custom Model Import\n", + "\n", + "This notebook shows how to use [Amazon Bedrock structured output](https://docs.aws.amazon.com/bedrock/latest/userguide/custom-model-import-advanced-features.html#structured-outputs) to constrain a model's output (both foundation and fine-tuned) to a JSON schema. The schema is enforced at the token level during decoding, so every response is guaranteed to parse and conform to your schema — no regex extraction, no malformed-output retries.\n", + "\n", + "\n", + "**Use case:** extracting structured metadata (entities, topics, sentiment, etc.) from a news article. The same pattern applies to any extraction task — invoices, support tickets, resumes, research abstracts.\n", + "\n", + "**Pipeline:**\n", + "```\n", + "article text → CMI model + json_schema constraint → validated JSON\n", + "```\n", + "\n", + "**Model & API format:** this notebook runs against `qwen.qwen3-coder-30b-a3b-v1:0` for easy demonstration, called via `InvokeModel` using the `OpenAIChatCompletion` request/response format (`{messages, max_tokens, response_format}`). \n", + "\n", + "CMI also supports `OpenAICompletion` and `BedrockCompletion` formats — see [Invoke an imported model](https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html) for the differences and [Structured outputs for imported models](https://docs.aws.amazon.com/bedrock/latest/userguide/custom-model-import-advanced-features.html#structured-outputs) for how `response_format` varies between them. \n", + "\n", + "To use your own fine-tuned model, replace `MODEL_ARN` with your CMI model ARN — as long as it accepts the `OpenAIChatCompletion` shape, no other changes are required.\n", + "\n", + "**References:**\n", + "1. [Invoke an imported model (request/response formats)](https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html)\n", + "2. [Structured outputs for Bedrock imported models](https://docs.aws.amazon.com/bedrock/latest/userguide/custom-model-import-advanced-features.html#structured-outputs)\n", + "3. [Get validated JSON results from models](https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html)\n", + "4. [Structured outputs on Amazon Bedrock: Schema-compliant AI responses](https://aws.amazon.com/blogs/machine-learning/structured-outputs-on-amazon-bedrock-schema-compliant-ai-responses/)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Setup" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Note: you may need to restart the kernel to use updated packages.\n" + ] + } + ], + "source": [ + "%pip install -q boto3 pydantic" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import time\n", + "from typing import Literal\n", + "\n", + "import boto3\n", + "from botocore.config import Config as BotoConfig\n", + "from pydantic import BaseModel, Field\n", + "\n", + "# !!!!!!!! REPLACE WITH YOUR CMI MODEL ARN !!!!!!!!\n", + "# i.e.: MODEL_ARN = \"arn:aws:bedrock:us-west-2::imported-model/\"\n", + "\n", + "# For demo purposes we use a foundation model so you can run the notebook — the API shape is identical.\n", + "MODEL_ARN = \"qwen.qwen3-coder-30b-a3b-v1:0\" # \"\" \n", + "REGION = \"us-east-1\"\n", + "\n", + "session = boto3.Session()\n", + "client = session.client(\n", + " \"bedrock-runtime\",\n", + " region_name=REGION,\n", + " config=BotoConfig(read_timeout=300, retries={\"max_attempts\": 0}),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "def invoke(payload, max_retries=6):\n", + " \"\"\"Call invoke_model with retry on transient errors.\n", + "\n", + " Imported CMI models may cold-start; ModelNotReady errors are retried with a longer backoff.\n", + " \"\"\"\n", + " last = None\n", + " for attempt in range(max_retries):\n", + " try:\n", + " resp = client.invoke_model(\n", + " modelId=MODEL_ARN,\n", + " body=json.dumps(payload),\n", + " contentType=\"application/json\",\n", + " accept=\"application/json\",\n", + " )\n", + " return json.loads(resp[\"body\"].read())\n", + " except Exception as e:\n", + " last = e\n", + " msg = str(e)\n", + " wait = 30 if \"ModelNotReady\" in msg else 10\n", + " print(f\" attempt {attempt+1}: {msg[:100]}... retrying in {wait}s\")\n", + " time.sleep(wait)\n", + " raise RuntimeError(f\"failed after {max_retries} attempts: {last}\")\n", + "\n", + "\n", + "def extract_content(result):\n", + " \"\"\"Return assistant text from an OpenAIChatCompletion response.\n", + "\n", + " Reasoning/CoT models may emit output in `reasoning` or `reasoning_content`\n", + " instead of (or in addition to) `content`. Fall back to those if `content` is empty.\n", + " \"\"\"\n", + " msg = result[\"choices\"][0][\"message\"]\n", + " text = msg.get(\"content\") or \"\"\n", + " if not text:\n", + " reasoning = msg.get(\"reasoning\") or msg.get(\"reasoning_content\") or \"\"\n", + " if reasoning and \"{\" in reasoning:\n", + " text = reasoning\n", + " return text" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Example input\n", + "\n", + "A short fictional news article. In a real pipeline this could come from an RSS feed, a CMS, or a scraped page." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "ARTICLE = \"\"\"Acme Robotics announced today that it has raised $120 million in a Series C funding round\n", + "led by Northwind Ventures, with participation from existing investors Pinegrove Capital and Helios Partners.\n", + "The Boston-based company, founded in 2018 by CEO Maria Chen and CTO David Okonkwo, plans to use the funding\n", + "to expand its warehouse automation business into the European market and double its engineering team by\n", + "the end of 2026.\n", + "\n", + "\"This round gives us the runway to bring our autonomous picking systems to retailers across the EU,\"\n", + "said Chen in a statement. \"Demand from logistics operators has tripled in the last 18 months.\"\n", + "\n", + "The deal, which closed on May 12, 2026, values Acme at approximately $850 million. The company\n", + "previously raised $45 million in a Series B led by Pinegrove in 2023.\"\"\"\n", + "\n", + "SYSTEM_PROMPT = \"\"\"You are an information extraction system. You read news articles and extract\n", + "structured metadata: headline, summary, entities (people, organizations, locations), topics,\n", + "sentiment, and any monetary or date values mentioned.\n", + "\n", + "Respond with a single JSON object only. Do not include prose, markdown fences, or commentary.\"\"\"\n", + "\n", + "USER_PROMPT = f\"\"\"Extract structured metadata from the following article and return it as JSON:\n", + "\n", + "{ARTICLE}\"\"\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Baseline — prompted JSON (no schema constraint)\n", + "\n", + "The prompt asks for JSON, but nothing forces it. Models often comply *most* of the time — and then drift into prose, markdown fences, extra commentary, or subtly malformed JSON whenever the input is unfamiliar or the temperature is high. Downstream code typically has to regex-extract a JSON block and hope it parses." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Baseline done in 3.4s\n", + "\n", + "{\n", + " \"headline\": \"Acme Robotics Raises $120 Million in Series C Funding\",\n", + " \"summary\": \"Acme Robotics secured $120 million in Series C funding led by Northwind Ventures, with participation from Pinegrove Capital and Helios Partners. The Boston-based company plans to expand into the European market and double its engineering team by 2026.\",\n", + " \"entities\": {\n", + " \"people\": [\"Maria Chen\", \"David Okonkwo\"],\n", + " \"organizations\": [\"Acme Robotics\", \"Northwind Ventures\", \"Pinegrove Capital\", \"Helios Partners\"],\n", + " \"locations\": [\"Boston\", \"EU\"]\n", + " },\n", + " \"topics\": [\"funding\", \"robotics\", \"warehouse automation\", \"expansion\"],\n", + " \"sentiment\": \"positive\",\n", + " \"monetary_values\": [120000000, 850000000, 45000000],\n", + " \"date_values\": [\"2026-05-12\", \"2023\", \"2026\"]\n", + "}\n" + ] + } + ], + "source": [ + "baseline_payload = {\n", + " \"messages\": [\n", + " {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n", + " {\"role\": \"user\", \"content\": USER_PROMPT},\n", + " ],\n", + " \"max_tokens\": 4096,\n", + " \"temperature\": 0.7,\n", + "}\n", + "\n", + "t0 = time.time()\n", + "baseline_result = invoke(baseline_payload)\n", + "print(f\"Baseline done in {time.time()-t0:.1f}s\\n\")\n", + "print(extract_content(baseline_result))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Define the output as a Pydantic schema\n", + "\n", + "Constrained decoding requires a JSON schema. Pydantic generates one automatically via `model_json_schema()`, which gives us a single source of truth for both the constraint and downstream validation." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"name\": \"Article\",\n", + " \"schema\": {\n", + " \"$defs\": {\n", + " \"Entities\": {\n", + " \"properties\": {\n", + " \"people\": {\n", + " \"items\": {\n", + " \"$ref\": \"#/$defs/Person\"\n", + " },\n", + " \"title\": \"People\",\n", + " \"type\": \"array\"\n", + " },\n", + " \"organizations\": {\n", + " \"items\": {\n", + " \"type\": \"string\"\n", + " },\n", + " \"title\": \"Organizations\",\n", + " \"type\": \"array\"\n", + " },\n", + " \"locations\": {\n", + " \"items\": {\n", + " \"type\": \"string\"\n", + " },\n", + " \"title\": \"Locations\",\n", + " \"type\": \"array\"\n", + " }\n", + " },\n", + " \"required\": [\n", + " \"people\",\n", + " \"organizations\",\n", + " \"locations\"\n", + " ],\n", + " \"title\": \"Entities\",\n", + " \"type\": \"object\"\n", + " },\n", + " \"MonetaryValue\": {\n", + " \"properties\": {\n", + " \"amount\": {\n", + " \"description\": \"Numeric amount\",\n", + " \"title\": \"Amount\",\n", + " \"type\": \"number\"\n", + " },\n", + " \"currency\": {\n", + " \"description\": \"ISO 4217 currency code, e.g. USD, EUR\",\n", + " \"title\": \"Currency\",\n", + " \"type\": \"string\"\n", + " },\n", + " \"context\": {\n", + " \"description\": \"What ...\n" + ] + } + ], + "source": [ + "class MonetaryValue(BaseModel):\n", + " amount: float = Field(description=\"Numeric amount\")\n", + " currency: str = Field(description=\"ISO 4217 currency code, e.g. USD, EUR\")\n", + " context: str = Field(description=\"What the amount refers to, e.g. 'Series C funding'\")\n", + "\n", + "\n", + "class Person(BaseModel):\n", + " name: str\n", + " role: str | None = Field(default=None, description=\"Title or role if mentioned, e.g. 'CEO'\")\n", + " organization: str | None = Field(default=None, description=\"Affiliated organization if mentioned\")\n", + "\n", + "\n", + "class Entities(BaseModel):\n", + " people: list[Person]\n", + " organizations: list[str]\n", + " locations: list[str]\n", + "\n", + "\n", + "class Article(BaseModel):\n", + " headline: str = Field(max_length=200, description=\"Concise headline summarizing the article\")\n", + " summary: str = Field(description=\"Two-to-three sentence summary of the article\")\n", + " category: Literal[\n", + " \"business\", \"technology\", \"politics\", \"sports\", \"science\", \"health\", \"entertainment\", \"other\"\n", + " ]\n", + " sentiment: Literal[\"positive\", \"negative\", \"neutral\", \"mixed\"]\n", + " topics: list[str] = Field(description=\"Short topic keywords, e.g. 'funding', 'robotics'\")\n", + " entities: Entities\n", + " monetary_values: list[MonetaryValue]\n", + " dates_mentioned: list[str] = Field(description=\"Dates referenced in the article, ISO 8601 if possible\")\n", + "\n", + "\n", + "article_schema = {\"name\": \"Article\", \"schema\": Article.model_json_schema()}\n", + "print(json.dumps(article_schema, indent=2)[:1200], \"...\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Generate with structured output\n", + "\n", + "Pass the schema via `response_format`. The model is now constrained to emit valid JSON matching the schema." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Structured output done in 2.2s\n", + "\n", + "{\n", + " \"headline\": \"Acme Robotics Raises $120 Million in Series C Funding\",\n", + " \"summary\": \"Acme Robotics secured $120 million in Series C funding led by Northwind Ventures, with participation from Pinegrove Capital and Helios Partners. The Boston-based company plans to expand into the European market and double its engineering team by 2026.\",\n", + " \"category\": \"technology\",\n", + " \"sentiment\": \"positive\",\n", + " \"topics\": [\"funding\", \"robotics\", \"expansion\", \"venture capital\"],\n", + " \"entities\": {\n", + " \"people\": [\n", + " {\"name\": \"Maria Chen\", \"role\": \"CEO\"},\n", + " {\"name\": \"David Okonkwo\", \"role\": \"CTO\"}\n", + " ],\n", + " \"organizations\": [\n", + " \"Acme Robotics\",\n", + " \"Northwind Ventures\",\n", + " \"Pinegrove Capital\",\n", + " \"Helios Partners\"\n", + " ],\n", + " \"locations\": [\n", + " \"Boston\",\n", + " \"EU\"\n", + " ]\n", + " },\n", + " \"monetary_values\": [\n", + " {\"amount\": 120000000, \"currency\": \"USD\", \"context\": \"Series C funding\"},\n", + " {\"amount\": 850000000, \"currency\": \"USD\", \"context\": \"company valuation\"},\n", + " {\"amount\": 45000000, \"currency\": \"USD\", \"context\": \"Series B funding\"}\n", + " ],\n", + " \"dates_mentioned\": [\n", + " \"2026-05-12\",\n", + " \"2023\",\n", + " \"2018\"\n", + " ]\n", + "}\n" + ] + } + ], + "source": [ + "structured_payload = {\n", + " \"messages\": [\n", + " {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n", + " {\"role\": \"user\", \"content\": USER_PROMPT},\n", + " ],\n", + " \"max_tokens\": 4096,\n", + " \"temperature\": 0.3,\n", + " \"response_format\": {\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": article_schema,\n", + " },\n", + "}\n", + "\n", + "t0 = time.time()\n", + "structured_result = invoke(structured_payload)\n", + "print(f\"Structured output done in {time.time()-t0:.1f}s\\n\")\n", + "\n", + "structured_json = extract_content(structured_result)\n", + "print(structured_json)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Validation\n", + "\n", + "Sanity checks: the output is valid JSON and matches the Pydantic schema." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "JSON parse: OK\n", + "Schema validation: OK\n", + "\n", + "Headline: Acme Robotics Raises $120 Million in Series C Funding\n", + "Category: technology\n", + "Sentiment: positive\n", + "People: ['Maria Chen', 'David Okonkwo']\n", + "Orgs: ['Acme Robotics', 'Northwind Ventures', 'Pinegrove Capital', 'Helios Partners']\n" + ] + } + ], + "source": [ + "parsed = json.loads(structured_json)\n", + "print(\"JSON parse: OK\")\n", + "\n", + "article = Article.model_validate(parsed)\n", + "print(\"Schema validation: OK\")\n", + "\n", + "print(f\"\\nHeadline: {article.headline}\")\n", + "print(f\"Category: {article.category}\")\n", + "print(f\"Sentiment: {article.sentiment}\")\n", + "print(f\"People: {[p.name for p in article.entities.people]}\")\n", + "print(f\"Orgs: {article.entities.organizations}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Automatic CMI performance improvements\n", + "\n", + "Per the [CMI performance enhancements announcement](https://aws.amazon.com/blogs/machine-learning/enhanced-performance-for-amazon-bedrock-custom-model-import/) (Nov 2025), existing imported models automatically benefit from PyTorch compilation + CUDA graph optimizations with no re-import required. Reported gains on reference models: TTFT reduced 76–88%, end-to-end latency reduced 18–59%, throughput up 25–29%. There is a one-time longer cold start on first deploy while artifacts compile and cache, then every scale-up after is faster than before." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".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.11.14" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +}