Update train.py and predict.py with new preprocessing and prediction logic

Author: mtandrita

Diabetes Prediction [END 2 END]/diabetes_pipeline/predict.py

@@ -1,30 +1,45 @@
+import argparse
 import joblib
-import numpy as np
-from config import MODEL_PATH, SCALER_PATH
-
-class DiabetesPredictor:
-	def __init__(self):
-		self.model = joblib.load(MODEL_PATH)
-		self.scaler = joblib.load(SCALER_PATH)
-
-	def predict(self, features: list) -> int:
-		"""
-		features order:
-		[Pregnancies, Glucose, BloodPressure, SkinThickness,
-		 Insulin, BMI, DPF, Age]
-		"""
-		features = np.array(features).reshape(1, -1)
-		features = self.scaler.transform(features)
-		return int(self.model.predict(features)[0])
-
-
-if __name__ == "__main__":
-	predictor = DiabetesPredictor()
-
-	sample_input = [2, 81, 72, 15, 76, 30.1, 0.547, 25]
-	result = predictor.predict(sample_input)
-
-	if result == 1:
-		print("Oops! You have diabetes.")
-	else:
-		print("Great! You don't have diabetes.")
+import pandas as pd
+
+MODEL_PATH = "model/diabetes_model.pkl"
+SCALER_PATH = "model/scaler.pkl"
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--pregnancies", type=int, required=True)
+parser.add_argument("--glucose", type=float, required=True)
+parser.add_argument("--bp", type=float, required=True)
+parser.add_argument("--skin", type=float, required=True)
+parser.add_argument("--insulin", type=float, required=True)
+parser.add_argument("--bmi", type=float, required=True)
+parser.add_argument("--dpf", type=float, required=True)
+parser.add_argument("--age", type=int, required=True)
+
+args = parser.parse_args()
+
+# Load model & scaler
+model = joblib.load(MODEL_PATH)
+scaler = joblib.load(SCALER_PATH)
+
+# IMPORTANT: feature names must match training
+input_data = pd.DataFrame([{
+	"Pregnancies": args.pregnancies,
+	"Glucose": args.glucose,
+	"BloodPressure": args.bp,
+	"SkinThickness": args.skin,
+	"Insulin": args.insulin,
+	"BMI": args.bmi,
+	"DPF": args.dpf,
+	"Age": args.age
+}])
+
+# Scale & predict
+input_scaled = scaler.transform(input_data)
+prediction = model.predict(input_scaled)[0]
+
+if prediction == 1:
+	print("⚠️ Diabetes detected")
+else:
+	print("✅ No diabetes detected")
+
+

Diabetes Prediction [END 2 END]/diabetes_pipeline/train.py

@@ -1,22 +1,29 @@
-# diabetes_pipeline/train.py
-
+import logging
 import joblib
 from sklearn.ensemble import RandomForestClassifier
 from data_preprocessing import load_and_preprocess
-from config import MODEL_DIR, MODEL_PATH, SCALER_PATH
+from config import MODEL_PATH, SCALER_PATH, MODEL_DIR
 
-# Ensure model directory exists
-MODEL_DIR.mkdir(exist_ok=True)
+# Logging setup
+logging.basicConfig(
+	filename="logs/training.log",
+	level=logging.INFO,
+	format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+logging.info("Training started")
 
-# Load and preprocess data
+# Load data
 X_train, X_test, y_train, y_test, scaler = load_and_preprocess()
 
 # Train model
 classifier = RandomForestClassifier(n_estimators=20, random_state=0)
 classifier.fit(X_train, y_train)
 
-# Save model and scaler
+# Save artifacts
+MODEL_DIR.mkdir(exist_ok=True)
 joblib.dump(classifier, MODEL_PATH)
 joblib.dump(scaler, SCALER_PATH)
 
-print("Model and scaler saved successfully.")
+logging.info("Model and scaler saved successfully")
+