02_ner/main.py

import argparse, torch
from transformers import pipeline

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--text", type=str, default="Barack Obama was born in Hawaii.")
    args = parser.parse_args()

    device = 0 if torch.cuda.is_available() else -1
    ner = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple", device=device)
    for ent in ner(args.text):
        print(ent)

if __name__ == "__main__":
    main()