wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
383 lines
13 KiB
Plaintext
383 lines
13 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Structured and Constrained LLM Output with Ludwig\n",
|
|
"\n",
|
|
"[](https://colab.research.google.com/github/ludwig-ai/ludwig/blob/main/examples/llm_structured_output/structured_output.ipynb)\n",
|
|
"\n",
|
|
"Large language models are trained to produce fluent text, but they have no built-in guarantee that their output follows a particular format. **Constrained decoding** solves this by restricting token sampling at inference time so the model can only ever produce output that satisfies a given constraint.\n",
|
|
"\n",
|
|
"**This notebook covers:**\n",
|
|
"1. Entity extraction with a JSON schema constraint\n",
|
|
"2. Sentiment classification with a regex constraint (guaranteed to produce only `positive`, `negative`, or `neutral`)\n",
|
|
"3. Logits extraction from LLM output\n",
|
|
"4. Side-by-side comparison: constrained vs unconstrained decoding\n",
|
|
"\n",
|
|
"All examples use small, freely available models that run on a free Colab T4 GPU."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!pip install \"ludwig[llm]\" --quiet"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"import textwrap\n",
|
|
"\n",
|
|
"import pandas as pd\n",
|
|
"import torch\n",
|
|
"import yaml\n",
|
|
"\n",
|
|
"from ludwig.api import LudwigModel\n",
|
|
"\n",
|
|
"# Check GPU availability\n",
|
|
"if torch.cuda.is_available():\n",
|
|
" device_name = torch.cuda.get_device_name(0)\n",
|
|
" vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
|
|
" print(f\"GPU: {device_name} ({vram_gb:.1f} GB VRAM)\")\n",
|
|
"else:\n",
|
|
" print(\"No GPU detected. Running on CPU — inference will be slower.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Entity extraction with JSON schema\n",
|
|
"\n",
|
|
"We configure Ludwig to constrain the LLM's output to a specific JSON schema. The schema describes the shape of the expected output — Ludwig compiles it into logit masks so only valid JSON tokens can be sampled.\n",
|
|
"\n",
|
|
"The model is `microsoft/phi-2` (2.7 B parameters, fits comfortably on a T4 with 4-bit quantization)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"entity_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: microsoft/phi-2\n",
|
|
"\n",
|
|
"quantization:\n",
|
|
" bits: 4\n",
|
|
"\n",
|
|
"prompt:\n",
|
|
" task: >\n",
|
|
" Extract the named entities from the input text and return them as a JSON\n",
|
|
" object with this structure:\n",
|
|
" {\"entities\": [{\"text\": \"...\", \"type\": \"PERSON|ORG|LOC|DATE\"}]}.\n",
|
|
" Return only valid JSON, nothing else.\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: text\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: output\n",
|
|
" type: text\n",
|
|
" decoder:\n",
|
|
" type: text_parser\n",
|
|
" json_schema:\n",
|
|
" type: object\n",
|
|
" properties:\n",
|
|
" entities:\n",
|
|
" type: array\n",
|
|
" items:\n",
|
|
" type: object\n",
|
|
" properties:\n",
|
|
" text:\n",
|
|
" type: string\n",
|
|
" type:\n",
|
|
" type: string\n",
|
|
" enum: [PERSON, ORG, LOC, DATE]\n",
|
|
" required: [text, type]\n",
|
|
" required: [entities]\n",
|
|
" additionalProperties: false\n",
|
|
"\n",
|
|
"generation:\n",
|
|
" max_new_tokens: 200\n",
|
|
" temperature: 0.1\n",
|
|
" do_sample: false\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"entity_samples = [\n",
|
|
" \"Apple Inc. was founded by Steve Jobs in Cupertino, California on April 1, 1976.\",\n",
|
|
" \"Elon Musk announced that Tesla will open a new Gigafactory in Berlin next year.\",\n",
|
|
" \"The United Nations headquarters is located in New York City.\",\n",
|
|
"]\n",
|
|
"\n",
|
|
"entity_df = pd.DataFrame({\"text\": entity_samples})\n",
|
|
"entity_df"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"entity_model = LudwigModel(config=entity_config)\n",
|
|
"entity_preds, _, _ = entity_model.predict(dataset=entity_df)\n",
|
|
"\n",
|
|
"for text, pred in zip(entity_samples, entity_preds[\"output_predictions\"]):\n",
|
|
" print(f\"Input: {text}\")\n",
|
|
" try:\n",
|
|
" parsed = json.loads(pred)\n",
|
|
" entities = parsed.get(\"entities\", [])\n",
|
|
" print(f\"Output: {len(entities)} entities\")\n",
|
|
" for ent in entities:\n",
|
|
" print(f\" - '{ent['text']}' ({ent['type']})\")\n",
|
|
" except json.JSONDecodeError:\n",
|
|
" print(f\"Output (raw): {pred}\")\n",
|
|
" print()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"The output is guaranteed to be valid JSON matching the schema. The LLM cannot produce malformed JSON, extra prose, or entity types outside the allowed enum.\n",
|
|
"\n",
|
|
"## Classification with constrained tokens\n",
|
|
"\n",
|
|
"For classification tasks we can constrain the output to a regex that only allows the valid class labels. Here we restrict the model to `positive`, `negative`, or `neutral`.\n",
|
|
"\n",
|
|
"We use `Qwen/Qwen2-0.5B-Instruct` (0.5 B parameters) — small enough to run quickly even on CPU."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"sentiment_constrained_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: Qwen/Qwen2-0.5B-Instruct\n",
|
|
"\n",
|
|
"prompt:\n",
|
|
" task: >\n",
|
|
" Classify the sentiment of the following text.\n",
|
|
" Respond with exactly one word: positive, negative, or neutral.\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: text\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: sentiment\n",
|
|
" type: text\n",
|
|
" decoder:\n",
|
|
" type: text_parser\n",
|
|
" # Regex constraint — only these three tokens can be emitted.\n",
|
|
" regex: \"(positive|negative|neutral)\"\n",
|
|
"\n",
|
|
"generation:\n",
|
|
" max_new_tokens: 10\n",
|
|
" temperature: 0.0\n",
|
|
" do_sample: false\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"sentiment_samples = [\n",
|
|
" \"I absolutely loved this product! It exceeded all my expectations.\",\n",
|
|
" \"The service was terrible and the food was cold.\",\n",
|
|
" \"The movie was okay, nothing special.\",\n",
|
|
" \"This is the best laptop I have ever owned. Highly recommend.\",\n",
|
|
" \"I waited two hours and they still got my order wrong.\",\n",
|
|
" \"The weather today is neither good nor bad.\",\n",
|
|
"]\n",
|
|
"\n",
|
|
"sentiment_df = pd.DataFrame({\"text\": sentiment_samples})\n",
|
|
"\n",
|
|
"sentiment_model = LudwigModel(config=sentiment_constrained_config)\n",
|
|
"sentiment_preds, _, _ = sentiment_model.predict(dataset=sentiment_df)\n",
|
|
"\n",
|
|
"for text, label in zip(sentiment_samples, sentiment_preds[\"sentiment_predictions\"]):\n",
|
|
" print(f\"{label!s:<10} {text}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Every prediction is one of the three valid labels. No post-processing or error handling is required.\n",
|
|
"\n",
|
|
"## Logits extraction\n",
|
|
"\n",
|
|
"Ludwig can return the raw logits (pre-softmax scores over the vocabulary) for each generated token. This is useful for:\n",
|
|
"\n",
|
|
"- Computing token-level confidence scores\n",
|
|
"- Calibration and uncertainty estimation\n",
|
|
"- Analysing what the model \"considered\" at each step\n",
|
|
"- Downstream ensemble or reranking tasks"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"logits_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: Qwen/Qwen2-0.5B-Instruct\n",
|
|
"\n",
|
|
"prompt:\n",
|
|
" task: \"Answer with a single word.\"\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: text\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: response\n",
|
|
" type: text\n",
|
|
" # Request logits to be returned alongside the prediction.\n",
|
|
" output_logits: true\n",
|
|
"\n",
|
|
"generation:\n",
|
|
" max_new_tokens: 5\n",
|
|
" temperature: 0.0\n",
|
|
" do_sample: false\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"logits_model = LudwigModel(config=logits_config)\n",
|
|
"logits_df = pd.DataFrame({\"text\": [\"Is Python a programming language?\"]})\n",
|
|
"\n",
|
|
"logits_preds, output_df, _ = logits_model.predict(dataset=logits_df, collect_predictions=True)\n",
|
|
"\n",
|
|
"print(\"Prediction:\", logits_preds[\"response_predictions\"].iloc[0])\n",
|
|
"\n",
|
|
"if \"response_logits\" in output_df.columns:\n",
|
|
" import numpy as np\n",
|
|
"\n",
|
|
" logits = output_df[\"response_logits\"].iloc[0]\n",
|
|
" logits_arr = np.array(logits)\n",
|
|
" print(f\"Logits shape: {logits_arr.shape}\")\n",
|
|
" # Convert to probabilities for the first generated token\n",
|
|
" probs = np.exp(logits_arr[0]) / np.exp(logits_arr[0]).sum()\n",
|
|
" top5_idx = np.argsort(probs)[::-1][:5]\n",
|
|
" print(\"Top-5 token probabilities (first generated token):\")\n",
|
|
" for idx in top5_idx:\n",
|
|
" print(f\" token {idx}: {probs[idx]:.4f}\")\n",
|
|
"else:\n",
|
|
" print(\"Logits column not present in output.\")\n",
|
|
" print(\"Available columns:\", list(output_df.columns))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Comparison: constrained vs unconstrained\n",
|
|
"\n",
|
|
"The following cell runs the same sentiment classification task with and without the regex constraint, then prints both outputs side by side to show how constrained decoding eliminates invalid responses."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Unconstrained config — same prompt, no decoder constraint\n",
|
|
"sentiment_unconstrained_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: Qwen/Qwen2-0.5B-Instruct\n",
|
|
"\n",
|
|
"prompt:\n",
|
|
" task: >\n",
|
|
" Classify the sentiment of the following text.\n",
|
|
" Respond with exactly one word: positive, negative, or neutral.\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: text\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: sentiment\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"generation:\n",
|
|
" max_new_tokens: 30\n",
|
|
" temperature: 0.7\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"unconstrained_model = LudwigModel(config=sentiment_unconstrained_config)\n",
|
|
"unconstrained_preds, _, _ = unconstrained_model.predict(dataset=sentiment_df)\n",
|
|
"\n",
|
|
"# The constrained model was already run above\n",
|
|
"constrained_labels = sentiment_preds[\"sentiment_predictions\"].tolist()\n",
|
|
"unconstrained_labels = unconstrained_preds[\"sentiment_predictions\"].tolist()\n",
|
|
"\n",
|
|
"valid_labels = {\"positive\", \"negative\", \"neutral\"}\n",
|
|
"\n",
|
|
"print(f\"{'Input':<48} {'Unconstrained':<32} {'Constrained'}\")\n",
|
|
"print(\"-\" * 100)\n",
|
|
"for text, unc, con in zip(sentiment_samples, unconstrained_labels, constrained_labels):\n",
|
|
" short = textwrap.shorten(text, width=46)\n",
|
|
" unc_str = str(unc).strip()\n",
|
|
" # Highlight invalid outputs\n",
|
|
" flag = \" *** INVALID\" if unc_str.lower() not in valid_labels else \"\"\n",
|
|
" print(f\"{short:<48} {unc_str:<32} {con!s}{flag}\")\n",
|
|
"\n",
|
|
"n_invalid = sum(1 for u in unconstrained_labels if str(u).strip().lower() not in valid_labels)\n",
|
|
"print(f\"\\nUnconstrained — invalid outputs: {n_invalid}/{len(sentiment_samples)}\")\n",
|
|
"print(f\"Constrained — invalid outputs: 0/{len(sentiment_samples)} (guaranteed)\")"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"accelerator": "GPU",
|
|
"colab": {
|
|
"gpuType": "T4",
|
|
"provenance": []
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"name": "python",
|
|
"version": "3.12.0"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
}
|