SabarnaDeb's picture
Upload folder using huggingface_hub
583485f verified
import os
import joblib
import pandas as pd
import gradio as gr
from huggingface_hub import hf_hub_download
# 1) SET YOUR MODEL REPO HERE (Model Hub repo, not dataset repo)
MODEL_REPO = os.getenv("MODEL_REPO", "SabarnaDeb/Capstone_PredictiveMaintenance_Model")
MODEL_FILE = os.getenv("MODEL_FILE", "model.joblib")
# 2) Download model file from Hugging Face Model Hub
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type="model")
model = joblib.load(model_path)
# 3) Feature list must match your training columns
FEATURES = [
"engine_rpm",
"lub_oil_pressure",
"fuel_pressure",
"coolant_pressure",
"lub_oil_temperature",
"coolant_temperature"
]
def predict(engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure,
lub_oil_temperature, coolant_temperature):
# 4) Save inputs into a DataFrame
data = {
"engine_rpm": [engine_rpm],
"lub_oil_pressure": [lub_oil_pressure],
"fuel_pressure": [fuel_pressure],
"coolant_pressure": [coolant_pressure],
"lub_oil_temperature": [lub_oil_temperature],
"coolant_temperature": [coolant_temperature],
}
input_df = pd.DataFrame(data)
# 5) Predict
pred = model.predict(input_df[FEATURES])[0]
prob = None
if hasattr(model, "predict_proba"):
prob = float(model.predict_proba(input_df[FEATURES])[:, 1][0])
# 6) Business-friendly output
if int(pred) == 1:
msg = "⚠️ Maintenance Needed"
else:
msg = "✅ Normal Operation"
if prob is not None:
msg += f"\nConfidence (maintenance probability): {prob:.2f}"
return msg, input_df
demo = gr.Interface(
fn=predict,
inputs=[
gr.Number(label="Engine RPM"),
gr.Number(label="Lub Oil Pressure"),
gr.Number(label="Fuel Pressure"),
gr.Number(label="Coolant Pressure"),
gr.Number(label="Lub Oil Temperature"),
gr.Number(label="Coolant Temperature"),
],
outputs=[
gr.Textbox(label="Prediction Result"),
gr.Dataframe(label="Input Data (saved as DataFrame)")
],
title="Predictive Maintenance – Engine Health",
description="Enter engine sensor readings to predict whether maintenance is needed."
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860, share=True)