from fastapi import FastAPI
from pydantic import BaseModel
import pickle
import numpy as np

app = FastAPI()

# Load your model
with open("model.pkl", "rb") as f:
    model = pickle.load(f)

class InputData(BaseModel):
    features: list  # list of numbers

@app.post("/predict")
def predict(data: InputData):
    X = np.array([data.features])

    # Get probability predictions (works for most sklearn models)
    probs = model.predict_proba(X)[0]

    return {
        "prob_0": float(probs[0]),
        "prob_1": float(probs[1])
    }

@app.get("/")
def root():
    return {"message": "Model API is working!"}