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}