| 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() | |