Spaces:
Running
Running
from fastapi import FastAPI | |
# from core.init_nlp import initialize_nlp | |
from contextlib import asynccontextmanager | |
import logging | |
from api.endpoints import location | |
from transformers import AutoTokenizer, AutoModelForTokenClassification | |
from transformers import pipeline | |
from models.fastapi_globals import g ,GlobalsMiddleware | |
async def lifespan(app: FastAPI): | |
# initialize_nlp() | |
print("Initializing NER model and tokenizer") | |
logging.info("Initializing NER model and tokenizer") | |
app.tokenizer = AutoTokenizer.from_pretrained("ml6team/bert-base-uncased-city-country-ner") | |
app.model = AutoModelForTokenClassification.from_pretrained("ml6team/bert-base-uncased-city-country-ner") | |
app.nlp = pipeline('ner', model=app.model, tokenizer=app.tokenizer, aggregation_strategy="simple") | |
g.set_default("ner_model", ner_model) | |
yield | |
del sentiment_model | |
g.cleanup() | |
app = FastAPI(lifespan=lifespan) | |
app.include_router(location.router, prefix="/location/api/v1") | |