magbert-ner-api / model.py
hassoudi's picture
Upload 6 files
d2237d8 verified
raw
history blame
425 Bytes
import os
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
HF_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = "TypicaAI/magbert-ner"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first")