from transformers import pipeline | |
import torch | |
ner_model = pipeline('ner') | |
model_checkpoint = "huggingface-course/bert-finetuned-ner" | |
classifier = pipeline("token-classification", model=model_checkpoint, aggregation_strategy="simple") | |
device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
def perform_ner(text): | |
# Your NER function implementation goes here | |
# Replace this with your own checkpoint | |
result = classifier(text) | |
return {"entities": [result]} | |