Spaces:
Sleeping
Sleeping
import os | |
#from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline | |
from transformers import pipeline | |
from huggingface_hub import login | |
# Safe writable cache dirs in Hugging Face Spaces move to space settings vars` | |
#os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache" | |
#os.environ["HF_HOME"] = "/tmp/hf_home" | |
#os.environ["HF_HUB_CACHE"] = "/tmp/hf_hub" | |
MODEL_ID = "TypicaAI/magbert-ner" | |
#HF_TOKEN = os.getenv("HF_TOKEN") | |
#tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN) | |
#model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, token=HF_TOKEN) | |
#ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first") | |
# Initialize global pipeline | |
ner_pipeline = None | |
# Authenticate using the secret `HFTOKEN` | |
def authenticate_with_token(): | |
"""Authenticate with the Hugging Face API using the HFTOKEN secret.""" | |
hf_token = os.getenv("HF_TOKEN") # Retrieve the token from environment variables | |
if not hf_token: | |
raise ValueError("HF_TOKEN is not set. Please add it to the Secrets in your Space settings.") | |
login(token=hf_token) | |
def load_healthcare_ner_pipeline(): | |
"""Load the Hugging Face pipeline for Healthcare NER.""" | |
global ner_pipeline | |
if ner_pipeline is None: | |
# Authenticate and initialize pipeline | |
authenticate_with_token() | |
ner_pipeline = pipeline( | |
"token-classification", | |
model=MODEL_ID, | |
aggregation_strategy="first" # Groups B- and I- tokens into entities | |
) | |
return ner_pipeline | |
# Get NER pipeline | |
ner_pipeline = load_healthcare_ner_pipeline() | |