1
# ==============================================================
2
# Author: Frees Ling
3
# Created: 2025-10-24
4
# Description: A script to train a simple logistic regression model
5
#              to predict student pass/fail based on hours studied.
6
# Version: 1.0
7
# ==============================================================
8
import argparse
9
import pandas as pd
10
import numpy as np
11
from sklearn.model_selection import train_test_split
12
from sklearn.preprocessing import StandardScaler
13
from sklearn.linear_model import LogisticRegression
14
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
15
import joblib
16
import os
17

18
from tensorflow.python.data.util.structure import NoneTensor
19

20

21
def load_data(csv_path):
22
    """
23
    CSV format: hours_studied, passed
24
    passed: 0 or 1
25
    """
26
    df = pd.read_csv(csv_path)
27
    # check
28
    assert 'hours_studied' in df.columns and 'passed' in df.columns, "CSV must have hours_studied and passed columns"
29
    X = df[['hours_studied']].values
30
    Y = df['passed'].values
31
    return X, Y
32

33
def train_save(csv_path, model_out='model.joblib', scaler_out='scaler.joblib', test_size=0.2, random_state=42):
34
    X, Y = load_data(csv_path)
35
    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=test_size, random_state=random_state, stratify=Y if len(np.unique(Y)) > 1 else None)
36

37
    scaler = StandardScaler()
38
    X_train_s = scaler.fit_transform(X_train)
39
    X_test_s = scaler.transform(X_test)
40

41
    model = LogisticRegression(max_iter=200)
42
    model.fit(X_train_s, Y_train)
43

44
    Y_pred = model.predict(X_test_s)
45
    print("Accuracy:", accuracy_score(Y_test, Y_pred))
46
    print("Confusion Matrix:\n", confusion_matrix(Y_test, Y_pred))
47
    print("Classification Report:\n", classification_report(Y_test, Y_pred))
48

49
    joblib.dump(model, model_out)
50
    joblib.dump(scaler, scaler_out)
51
    print(f"Saved model -> {model_out}, scaler -> {scaler_out}")
52

53
if __name__ == "__main__":
54
    parser = argparse.ArgumentParser(description='Simple AI First test')
55
    parser.add_argument('--csv', type=str, default="data.csv", help='Path to the CSV data file')
56
    parser.add_argument('--model_out', type=str, default="model.joblib", help='Path to save the trained model')
57
    parser.add_argument('--scaler_out', type=str, default="scaler.joblib", help='Path to save the scaler')
58
    parser.add_argument('--test_size', type=float, default=0.2, help='Test size')
59
    parser.add_argument('--random_state', type=int, default=42, help='Random state for reproducibility')
60
    args = parser.parse_args()
61

62
    if not os.path.exists(args.csv):
63
        #If no csv file, create a sample one
64
        sample = pd.DataFrame({
65
            'hours_studied': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
66
            'passed': [0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
67
        })
68
        sample.to_csv(args.csv, index=False)
69
        print(f"Sample data created at {args.csv}. Edit it to add your data, then re-run training.")
70
    train_save(args.csv, args.model_out, args.scaler_out, args.test_size, args.random_state)