Whisper / ner.py
jaimin's picture
Upload 2 files
2f3b04f verified
raw
history blame
477 Bytes
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]}