train.py

import os
import pickle
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import Pipeline
from sklearn.metrics import accuracy_score

def train():
    # load dataset
    iris = load_iris(as_frame=True)
    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)
    
    # pipeline preprocessing + model
    clf = Pipeline([
        ("scaler", StandardScaler()),
        ("rf", RandomForestClassifier(n_estimators=100, random_state=42))
    ])
    
    # train
    clf.fit(X_train, y_train)
    
    # evaluate
    preds = clf.predict(X_test)
    acc = accuracy_score(y_test, preds)
    print(f"Accuracy: {acc:.3f}")
    
    # make sure models/ folder exists
    os.makedirs("models", exist_ok=True)
    
    # save model
    with open("models/model.pkl", "wb") as f:
        pickle.dump(clf, f)

if __name__ == "__main__":
    train()