app.py

from huggingface_hub import login
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import os

app = FastAPI()

access_token = os.environ.get("ACCESS_TOKEN_1")
login(token=access_token, add_to_git_credential=True)

# Load the model and tokenizer from the Hugging Face Hub
model_name = "MHULO/yembaner"
nlp = pipeline("ner", model=model_name, tokenizer=model_name)

class TextRequest(BaseModel):
    text: str

@app.post("/predict/")
def predict(request: TextRequest):
    ner_results = nlp(request.text)
    return ner_results


@app.get("/")
def root():
    return {"prediction url": "/predict/"}