Update app.py

app.py

@@ -38,4 +38,46 @@ def train_model():
     model.fit(X_train, y_train)
 
     joblib.dump(model, MODEL_PATH)
-    print("Model sav
+    print("Model saved as rf_model.pkl")
+    return model
+
+
+# Load or train model
+if os.path.exists(MODEL_PATH):
+    print("Loading existing model...")
+    model = joblib.load(MODEL_PATH)
+else:
+    model = train_model()
+
+
+# ---------------------------
+# PREDICTION FUNCTION
+# ---------------------------
+
+feature_names = [
+    'fixed acidity', 'volatile acidity', 'citric acid', 'residual sugar',
+    'chlorides', 'free sulfur dioxide', 'total sulfur dioxide', 'density',
+    'pH', 'sulphates', 'alcohol'
+]
+
+def predict_quality(*inputs):
+    df = pd.DataFrame([inputs], columns=feature_names)
+    prediction = model.predict(df)[0]
+    return f"Predicted Wine Quality: {prediction}"
+
+
+# ---------------------------
+# GRADIO UI
+# ---------------------------
+
+inputs_ui = [gr.Number(label=name) for name in feature_names]
+
+demo = gr.Interface(
+    fn=predict_quality,
+    inputs=inputs_ui,
+    outputs=gr.Textbox(label="Prediction"),
+    title="🍾 White Wine Quality Predictor (Trains on HF Space)",
+    description="Random Forest model trained on the UCI White Wine Quality dataset."
+)
+
+demo.launch()