ml_api / app.py
murfhi's picture
Create app.py
24e1f4d verified
raw
history blame contribute delete
726 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import numpy as np
app = FastAPI()
# Load model
model = joblib.load('BEST_MODEL_LightGBM_TFIDF.joblib')
class PredictionRequest(BaseModel):
text: str
top_k: int = 3
@app.post("/predict")
def predict(request: PredictionRequest):
# Your prediction logic here
probabilities = model.predict_proba([request.text])[0]
top_indices = np.argsort(probabilities)[::-1][:request.top_k]
results = []
for idx in top_indices:
results.append({
'sdg_number': idx + 1,
'sdg_name': f'SDG {idx + 1}',
'confidence': float(probabilities[idx])
})
return {'predictions': results}