train_model.py

# ============================================================
#  train_model.py — Train and save ML model
#  Run this ONCE to generate model.pkl
#  Author  : Aarica Raj
# ============================================================

from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import pickle

# ── Load Dataset ────────────────────────────────────────────
print("Loading Iris dataset...")
iris    = load_iris()
X, y    = iris.data, iris.target

# ── Split Data ──────────────────────────────────────────────
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# ── Train Model ─────────────────────────────────────────────
print("Training Random Forest model...")
model = RandomForestClassifier(
    n_estimators = 100,
    random_state = 42
)
model.fit(X_train, y_train)

# ── Evaluate ────────────────────────────────────────────────
y_pred   = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"Model Accuracy: {accuracy * 100:.2f}%")

# ── Save Model ──────────────────────────────────────────────
with open("model.pkl", "wb") as f:
    pickle.dump(model, f)

print("Model saved as model.pkl!")
print("Now run: docker build -t mlops-pipeline .")