Extract named entities from text

%pip install -qU pixeltable openai

import getpass
import os

if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')

import json
import pixeltable as pxt
from pixeltable.functions.openai import chat_completions

# Create a fresh directory
pxt.drop_dir('entities_demo', force=True)
pxt.create_dir('entities_demo')

# Define the JSON schema for entity extraction
entity_schema = {
    'type': 'json_schema',
    'json_schema': {
        'name': 'entities',
        'strict': True,
        'schema': {
            'type': 'object',
            'properties': {
                'people': {
                    'type': 'array',
                    'items': {'type': 'string'},
                    'description': 'Names of people mentioned',
                },
                'organizations': {
                    'type': 'array',
                    'items': {'type': 'string'},
                    'description': 'Names of companies, institutions, or groups',
                },
                'locations': {
                    'type': 'array',
                    'items': {'type': 'string'},
                    'description': 'Geographic locations (cities, countries, addresses)',
                },
                'dates': {
                    'type': 'array',
                    'items': {'type': 'string'},
                    'description': 'Dates or time references',
                },
            },
            'required': ['people', 'organizations', 'locations', 'dates'],
            'additionalProperties': False,
        },
    },
}

# Create table for articles
articles = pxt.create_table(
    'entities_demo/articles', {'title': pxt.String, 'content': pxt.String}
)

# Add entity extraction column
extraction_prompt = (
    'Extract all named entities from the following text:\n\n'
    + articles.content
)

articles.add_computed_column(
    extraction_response=chat_completions(
        messages=[{'role': 'user', 'content': extraction_prompt}],
        model='gpt-4o-mini',
        model_kwargs={'response_format': entity_schema},
    )
)

# Extract the entities JSON
articles.add_computed_column(
    entities=articles.extraction_response.choices[0].message.content
)

# Insert sample articles
sample_articles = [
    {
        'title': 'Tech Acquisition',
        'content': 'Microsoft announced today that CEO Satya Nadella will lead the acquisition of a Seattle-based startup. The deal, expected to close in March 2024, is valued at $500 million.',
    },
    {
        'title': 'Sports Update',
        'content': 'LeBron James led the Los Angeles Lakers to victory against the Boston Celtics on Tuesday night at Staples Center. Coach Darvin Ham praised the teams performance.',
    },
    {
        'title': 'Research Breakthrough',
        'content': 'Dr. Sarah Chen at Stanford University published groundbreaking research on renewable energy. The study, funded by the National Science Foundation, was conducted in Palo Alto, California.',
    },
]

articles.insert(sample_articles)

# View extracted entities
articles.select(articles.title, articles.entities).collect()