Simple classification model for iris data

iris_classification.py

def irisDataClassification():
    # Import modules
    from sklearn import datasets
    from sklearn.model_selection import train_test_split
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import accuracy_score

    # Import some data to play with
    iris = datasets.load_iris()
    X, y = iris.data, iris.target
    
    # Data splitting
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 1, stratify = y)
    
    # Create logistic regression object
    model = LogisticRegression()
    
    # Data modelling with logistic regression
    model.fit(X_train, y_train)
    
    # Create prediction using testing data
    y_pred = model.predict(X_test)
    
    # Print out the accuracy
    accuracy = accuracy_score(y_test, y_pred)
    print(accuracy)

if __name__ == '__main__':
    irisDataClassification()