File size: 477 Bytes
2f3b04f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
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]}
|