ART / app.py
ksrivarshith's picture
Create app.py
47214b0 verified
Raw
History Blame Contribute Delete
1.29 kB
from fastapi import FastAPI
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
app = FastAPI()
# ================= GLOBAL MODEL =================
model = None
scaler = None
features = None
# ================= TRAIN API =================
@app.post("/train")
def train_model(data: dict):
global model, scaler, features
df = pd.DataFrame(data["dataset"])
target = data["target"]
X = df.drop(columns=[target])
y = df[target]
features = list(X.columns)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
model = LogisticRegression(max_iter=5000)
model.fit(X_scaled, y)
return {"status": "Model trained successfully"}
# ================= PREDICT API =================
@app.post("/predict")
def predict(input_data: dict):
global model, scaler, features
if model is None:
return {"error": "Model not trained yet"}
df = pd.DataFrame([input_data])
# Ensure correct column order
df = df[features]
X_scaled = scaler.transform(df)
pred = model.predict(X_scaled)[0]
prob = float(np.max(model.predict_proba(X_scaled)))
return {
"prediction": int(pred),
"confidence": prob
}