Agent-37's picture
Fix app.py
96cf76a
Raw
History Blame Contribute Delete
2.06 kB
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import json
import pandas as pd
app = FastAPI(title="Game of Thrones House Predictor")
model = joblib.load("model.pkl")
with open("feature_columns.json", "r") as f:
feature_columns = json.load(f)["columns"]
with open("label_classes.json", "r") as f:
label_classes = json.load(f)["classes"]
class CharacterInput(BaseModel):
region: str
primary_role: str
alignment: str
status: str
species: str
honour_1to5: int
ruthlessness_1to5: int
intelligence_1to5: int
combat_skill_1to5: int
diplomacy_1to5: int
leadership_1to5: int
trait_loyal: bool
trait_scheming: bool
class PredictionOutput(BaseModel):
predicted_house: str
@app.get("/")
def root():
return {"message": "Game of Thrones House Predictor API", "docs": "/docs"}
@app.post("/predict", response_model=PredictionOutput)
def predict(character: CharacterInput):
input_dict = {
"region": character.region,
"primary_role": character.primary_role,
"alignment": character.alignment,
"status": character.status,
"species": character.species,
"honour_1to5": character.honour_1to5,
"ruthlessness_1to5": character.ruthlessness_1to5,
"intelligence_1to5": character.intelligence_1to5,
"combat_skill_1to5": character.combat_skill_1to5,
"diplomacy_1to5": character.diplomacy_1to5,
"leadership_1to5": character.leadership_1to5,
"trait_loyal": int(character.trait_loyal),
"trait_scheming": int(character.trait_scheming),
}
df = pd.DataFrame([input_dict])
df_encoded = pd.get_dummies(df, dummy_na=True)
for col in feature_columns:
if col not in df_encoded.columns:
df_encoded[col] = 0
df_encoded = df_encoded[feature_columns]
prediction = model.predict(df_encoded)[0]
return PredictionOutput(predicted_house=prediction)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)