Entity Extraction for Knowledge Graphs

What Is Entity Extraction?

spaCy NER Basics

import spacy

# Load English model (install: python -m spacy download en_core_web_sm)
nlp = spacy.load('en_core_web_sm')

text = 'Elon Musk founded SpaceX in 2002 in Hawthorne, California. Tesla is headquartered in Austin, Texas.'

doc = nlp(text)

for ent in doc.ents:
    print(f'{ent.text:30} {ent.label_:15} {spacy.explain(ent.label_)}')

# Output:
# Elon Musk                      PERSON          People, including fictional
# SpaceX                         ORG             Companies, agencies...
# 2002                           DATE            Absolute or relative dates
# Hawthorne, California          GPE             Countries, cities, states

Extracting Entities as Structured Data

import spacy
from collections import defaultdict

nlp = spacy.load('en_core_web_sm')

def extract_entities(text: str) -> dict:
    doc = nlp(text)
    entities = defaultdict(set)
    
    for ent in doc.ents:
        entities[ent.label_].add(ent.text.strip())
    
    # Convert sets to lists for JSON serialization
    return {k: list(v) for k, v in entities.items()}

text = 'Apple CEO Tim Cook announced new products. The event was held in Cupertino on September 12, 2023.'
result = extract_entities(text)

import json
print(json.dumps(result, indent=2))
# {
#   "ORG": ["Apple"],
#   "PERSON": ["Tim Cook"],
#   "GPE": ["Cupertino"],
#   "DATE": ["September 12, 2023"]
# }

Relation Extraction with LLM

import openai
import json

client = openai.OpenAI(api_key='sk-...')

def extract_relations(text: str, entities: list) -> list:
    if len(entities) < 2:
        return []
    
    entity_list = ', '.join(entities)
    prompt = (
        f'Given this text and these entities: {entity_list}\n\n'
        f'Text: {text}\n\n'
        'Extract relationships between the entities. '
        'Return a JSON array of objects with fields: '
        'subject (string), relation (string), object (string). '
        'Use concise relation labels like FOUNDED_BY, WORKS_AT, LOCATED_IN, ACQUIRED_BY.'
    )
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}],
        response_format={'type': 'json_object'}
    )
    result = json.loads(response.choices[0].message.content)
    return result.get('relations', [])

text = 'Elon Musk founded SpaceX in Hawthorne, California.'
entities = ['Elon Musk', 'SpaceX', 'Hawthorne']
relations = extract_relations(text, entities)
print(relations)
# [{'subject': 'Elon Musk', 'relation': 'FOUNDED', 'object': 'SpaceX'},
#  {'subject': 'SpaceX', 'relation': 'LOCATED_IN', 'object': 'Hawthorne'}]

Entity Normalization

import openai
import json

client = openai.OpenAI(api_key='sk-...')

def normalize_entities(entity_mentions: list, entity_type: str) -> dict:
    '''
    Groups variations of the same entity together.
    Returns {canonical_name: [mention1, mention2, ...]}
    '''
    if len(entity_mentions) <= 1:
        return {m: [m] for m in entity_mentions}
    
    mentions_str = json.dumps(entity_mentions)
    prompt = (
        f'These are {entity_type} entity mentions from text: {mentions_str}\n'
        'Group mentions that refer to the same entity. '
        'Return JSON: {"groups": [["canonical", "alias1", "alias2"], ...]}'
    )
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}],
        response_format={'type': 'json_object'}
    )
    result = json.loads(response.choices[0].message.content)
    
    normalized = {}
    for group in result.get('groups', []):
        if group:
            canonical = group[0]
            for mention in group:
                normalized[mention] = canonical
    return normalized

mentions = ['Elon Musk', 'Musk', 'E. Musk', 'Tim Cook']
result = normalize_entities(mentions, 'PERSON')
print(result)

Entity Deduplication

from difflib import SequenceMatcher

class EntityRegistry:
    def __init__(self, similarity_threshold=0.85):
        self.entities = {}  # {canonical_name: entity_data}
        self.threshold = similarity_threshold
    
    def similarity(self, a: str, b: str) -> float:
        return SequenceMatcher(None, a.lower(), b.lower()).ratio()
    
    def find_existing(self, name: str) -> str:
        for canonical in self.entities:
            if self.similarity(name, canonical) >= self.threshold:
                return canonical
        return None
    
    def add_entity(self, name: str, entity_type: str, metadata: dict = None) -> str:
        existing = self.find_existing(name)
        if existing:
            print(f'Merged "{name}" -> "{existing}"')
            return existing
        self.entities[name] = {'type': entity_type, 'metadata': metadata or {}}
        print(f'New entity: "{name}"')
        return name

registry = EntityRegistry()
registry.add_entity('OpenAI', 'ORG', {})
registry.add_entity('Open AI', 'ORG', {})   # Merged
registry.add_entity('OpenAI Inc', 'ORG', {})  # Merged
registry.add_entity('Microsoft', 'ORG', {})

Processing Multiple Documents

import spacy
from typing import List, Dict

nlp = spacy.load('en_core_web_sm')

@dataclass
class KGFact:
    subject: str
    relation: str
    obj: str
    source_doc: str

def process_corpus(documents: List[Dict]) -> List:
    facts = []
    entity_registry = EntityRegistry()
    
    for doc in documents:
        doc_id = doc['id']
        text = doc['text']
        
        # Extract entities
        spacy_doc = nlp(text)
        persons = [ent.text for ent in spacy_doc.ents if ent.label_ == 'PERSON']
        orgs = [ent.text for ent in spacy_doc.ents if ent.label_ == 'ORG']
        
        # Register entities (deduplication)
        for p in persons:
            entity_registry.add_entity(p, 'PERSON')
        for o in orgs:
            entity_registry.add_entity(o, 'ORG')
        
        print(f'Processed doc {doc_id}: {len(persons)} persons, {len(orgs)} orgs')
    
    return entity_registry.entities

docs = [
    {'id': 'doc1', 'text': 'Sundar Pichai leads Google.'},
    {'id': 'doc2', 'text': 'Google CEO Sundar Pichai spoke at the conference.'}
]
entities = process_corpus(docs)

Entity Types for Knowledge Graphs

ENTITY_TYPES = {
    'PERSON': {
        'description': 'Individual human being',
        'properties': ['name', 'title', 'email'],
        'spacy_labels': ['PERSON']
    },
    'ORGANIZATION': {
        'description': 'Company, institution, or group',
        'properties': ['name', 'industry', 'founded'],
        'spacy_labels': ['ORG']
    },
    'LOCATION': {
        'description': 'Geographic place',
        'properties': ['name', 'country', 'coordinates'],
        'spacy_labels': ['GPE', 'LOC', 'FAC']
    },
    'PRODUCT': {
        'description': 'Product or service (custom, not in spaCy)',
        'properties': ['name', 'category', 'version'],
        'spacy_labels': ['PRODUCT']
    }
}

# Map spaCy labels to our types
def map_spacy_to_entity_type(spacy_label: str) -> str:
    for entity_type, config in ENTITY_TYPES.items():
        if spacy_label in config['spacy_labels']:
            return entity_type
    return 'UNKNOWN'

print(map_spacy_to_entity_type('ORG'))     # ORGANIZATION
print(map_spacy_to_entity_type('GPE'))     # LOCATION
print(map_spacy_to_entity_type('PERSON'))  # PERSON

Combining spaCy and LLM

import spacy
import openai

nlp = spacy.load('en_core_web_sm')
client = openai.OpenAI(api_key='sk-...')

def full_extraction_pipeline(text: str) -> dict:
    # Step 1: Fast spaCy extraction
    doc = nlp(text)
    entities = [
        {'text': ent.text, 'type': ent.label_}
        for ent in doc.ents
        if ent.label_ in ['PERSON', 'ORG', 'GPE', 'PRODUCT']
    ]
    
    if len(entities) < 2:
        return {'entities': entities, 'relations': []}
    
    # Step 2: LLM relation extraction (only when we have multiple entities)
    entity_names = [e['text'] for e in entities]
    relations = extract_relations(text, entity_names)
    
    return {
        'entities': entities,
        'relations': relations
    }

text = 'Satya Nadella of Microsoft acquired Activision Blizzard for $69 billion.'
result = full_extraction_pipeline(text)
print('Entities:', result['entities'])
print('Relations:', result['relations'])

Validating Extracted Data

VALID_RELATIONS = {
    'FOUNDED_BY', 'WORKS_AT', 'LOCATED_IN', 'ACQUIRED_BY',
    'PARTNERED_WITH', 'INVESTED_IN', 'CEO_OF', 'PRODUCT_OF'
}

def validate_relation(relation: dict, known_entities: set) -> tuple:
    errors = []
    subject = relation.get('subject', '')
    relation_label = relation.get('relation', '')
    obj = relation.get('object', '')
    
    if not subject:
        errors.append('Missing subject')
    if not obj:
        errors.append('Missing object')
    if relation_label not in VALID_RELATIONS:
        errors.append(f'Unknown relation: {relation_label}')
    if subject and subject not in known_entities:
        errors.append(f'Unknown entity: {subject}')
    if obj and obj not in known_entities:
        errors.append(f'Unknown entity: {obj}')
    
    return len(errors) == 0, errors

known = {'Elon Musk', 'SpaceX', 'California'}
relation = {'subject': 'Elon Musk', 'relation': 'FOUNDED_BY', 'object': 'SpaceX'}
valid, errors = validate_relation(relation, known)
print('Valid:', valid, 'Errors:', errors)

Incremental Graph Building

def build_knowledge_graph_incremental(new_documents: list, graph_db, entity_registry):
    for doc in new_documents:
        print(f'Processing document: {doc["id"]}')
        
        # Extract
        extraction = full_extraction_pipeline(doc['text'])
        
        # Register and deduplicate entities
        canonical_entities = {}
        for entity in extraction['entities']:
            canonical = entity_registry.add_entity(
                entity['text'],
                entity['type']
            )
            canonical_entities[entity['text']] = canonical
            # Upsert node in graph
            graph_db.upsert_node(canonical, entity['type'])
        
        # Validate and insert relations
        for relation in extraction['relations']:
            # Map to canonical names
            subj = canonical_entities.get(relation['subject'], relation['subject'])
            obj = canonical_entities.get(relation['object'], relation['object'])
            valid, errors = validate_relation(
                {'subject': subj, 'relation': relation['relation'], 'object': obj},
                set(canonical_entities.values())
            )
            if valid:
                graph_db.upsert_edge(subj, relation['relation'], obj, doc['id'])
            else:
                print(f'Skipping invalid relation: {errors}')

print('Graph build pipeline defined')

Knowledge Check: Entity Extraction

Entity Extraction Summary