02_ner/inference.py

import torch
from transformers import pipeline

MODEL_NAME = "dslim/bert-base-NER"
device = 0 if torch.cuda.is_available() else -1

ner = pipeline(
    "ner",
    model=MODEL_NAME,
    aggregation_strategy="simple",
    device=device
)

texts = [
    "Barack Obama was born in Hawaii.",
    "Elon Musk founded SpaceX in California."
]

for text in texts:
    print(f"Text: {text}")
    for ent in ner(text):
        print(ent)
    print()