{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "header" }, "source": [ "# Romeo and Juliet Text Extraction with LangExtract\n", "\n", "This notebook demonstrates extracting characters, emotions, and relationships from Shakespeare's Romeo and Juliet using LangExtract.\n", "\n", "[](https://colab.research.google.com/github/google/langextract/blob/main/examples/notebooks/romeo_juliet_extraction.ipynb)" ] }, { "cell_type": "markdown", "metadata": { "id": "setup_header" }, "source": [ "## Setup" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "id": "install" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Note: you may need to restart the kernel to use updated packages.\n" ] } ], "source": [ "# Install LangExtract\n", "%pip install -q langextract" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "id": "api_key" }, "outputs": [], "source": [ "# Set up your Gemini API key\n", "# Get your key from: https://aistudio.google.com/app/apikey\n", "import os\n", "from getpass import getpass\n", "\n", "if 'GEMINI_API_KEY' not in os.environ:\n", " os.environ['GEMINI_API_KEY'] = getpass('Enter your Gemini API key: ')" ] }, { "cell_type": "markdown", "metadata": { "id": "define_header" }, "source": [ "## Define Extraction Task" ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "id": "setup_extraction" }, "outputs": [], "source": [ "import langextract as lx\n", "import textwrap\n", "\n", "# Define the extraction task\n", "prompt = textwrap.dedent(\"\"\"\\\n", " Extract characters, emotions, and relationships in order of appearance.\n", " Use exact text for extractions. Do not paraphrase or overlap entities.\n", " Provide meaningful attributes for each entity to add context.\"\"\")\n", "\n", "# Provide a high-quality example\n", "examples = [\n", " lx.data.ExampleData(\n", " text=\"ROMEO. But soft! What light through yonder window breaks? It is the east, and Juliet is the sun.\",\n", " extractions=[\n", " lx.data.Extraction(\n", " extraction_class=\"character\",\n", " extraction_text=\"ROMEO\",\n", " attributes={\"emotional_state\": \"wonder\"}\n", " ),\n", " lx.data.Extraction(\n", " extraction_class=\"emotion\",\n", " extraction_text=\"But soft!\",\n", " attributes={\"feeling\": \"gentle awe\"}\n", " ),\n", " lx.data.Extraction(\n", " extraction_class=\"relationship\",\n", " extraction_text=\"Juliet is the sun\",\n", " attributes={\"type\": \"metaphor\"}\n", " ),\n", " ]\n", " )\n", "]" ] }, { "cell_type": "markdown", "metadata": { "id": "extract_header" }, "source": [ "## Extract from Sample Text" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "id": "simple_extraction" }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[94m\u001b[1mLangExtract\u001b[0m: model=\u001b[92mgemini-2.5-flash\u001b[0m, current=\u001b[92m68\u001b[0m chars, processed=\u001b[92m68\u001b[0m chars: [00:01]" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\u001b[92m✓\u001b[0m Extraction processing complete\n", "\u001b[92m✓\u001b[0m Extracted \u001b[1m3\u001b[0m entities (\u001b[1m3\u001b[0m unique types)\n", " \u001b[96m•\u001b[0m Time: \u001b[1m1.96s\u001b[0m\n", " \u001b[96m•\u001b[0m Speed: \u001b[1m35\u001b[0m chars/sec\n", " \u001b[96m•\u001b[0m Chunks: \u001b[1m1\u001b[0m\n", "Extracted 3 entities:\n", "\n", "• character: 'Lady Juliet'\n", " - emotional_state: longing\n", "• emotion: 'gazed longingly at the stars, her heart aching'\n", " - feeling: melancholy longing\n", "• relationship: 'her heart aching for Romeo'\n", " - type: romantic\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "# Simple extraction from a short text\n", "input_text = \"Lady Juliet gazed longingly at the stars, her heart aching for Romeo\"\n", "\n", "result = lx.extract(\n", " text_or_documents=input_text,\n", " prompt_description=prompt,\n", " examples=examples,\n", " model_id=\"gemini-2.5-flash\",\n", ")\n", "\n", "# Display results\n", "print(f\"Extracted {len(result.extractions)} entities:\\n\")\n", "for extraction in result.extractions:\n", " print(f\"• {extraction.extraction_class}: '{extraction.extraction_text}'\")\n", " if extraction.attributes:\n", " for key, value in extraction.attributes.items():\n", " print(f\" - {key}: {value}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "viz_header" }, "source": [ "## Interactive Visualization" ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "id": "visualization" }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[94m\u001b[1mLangExtract\u001b[0m: Saving to \u001b[92mromeo_juliet.jsonl\u001b[0m: 1 docs [00:00, 995.33 docs/s]" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\u001b[92m✓\u001b[0m Saved \u001b[1m1\u001b[0m documents to \u001b[92mromeo_juliet.jsonl\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n", "\u001b[94m\u001b[1mLangExtract\u001b[0m: Loading \u001b[92mromeo_juliet.jsonl\u001b[0m: 100%|██████████| 961/961 [00:00<00:00, 2.49MB/s]" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\u001b[92m✓\u001b[0m Loaded \u001b[1m1\u001b[0m documents from \u001b[92mromeo_juliet.jsonl\u001b[0m\n", "Interactive visualization (hover over highlights to see attributes):\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n" ] }, { "data": { "text/html": [ "\n", "