mrgmd01's picture
Update app.py
79de231 verified
Raw
History Blame Contribute Delete
1.13 kB
import os
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import pipeline
app = FastAPI()
token = os.environ.get("HF_TOKEN")
english_model = None
urdu_model = None
def get_english_model():
global english_model
if english_model is None:
english_model = pipeline("text-classification", model="mrgmd01/Finetuned_siebert-sentiment-roberta-large-english", token=token)
return english_model
def get_urdu_model():
global urdu_model
if urdu_model is None:
urdu_model = pipeline("text-classification", model="mrgmd01/SA_Model_bert-base-multilingual-uncased", token=token)
return urdu_model
@app.get("/")
def root():
return {"status": "running"}
@app.post("/predict")
async def predict(request: Request):
body = await request.json()
text = body.get("text", "")
language = body.get("language", "english")
if language in ["urdu", "roman_urdu"]:
result = get_urdu_model()(text)[0]
else:
result = get_english_model()(text)[0]
return JSONResponse({"label": result["label"], "score": float(result["score"])})