app/download_ner.py

import os
from huggingface_hub import login
from transformers import pipeline

hf_token = os.getenv('HF_TOKEN')
if hf_token:
    login(token=hf_token)

cache_dir = '/app/hf_cache'
model_name = 'dslim/distilbert-NER'

expected_cache_path = os.path.join(cache_dir, 'transformers', model_name.replace('/', '--'))

if not os.path.exists(expected_cache_path):
    print(f'Downloading NER model: {model_name}...')
    nlp = pipeline('token-classification', model=model_name)
else:
    print(f'NER model ({model_name}) already cached. Skipping download.')

print('NER model download check complete.')