main.py

import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

iris = datasets.load_iris()

def plot():
    _, ax = plt.subplots()
    scatter = ax.scatter(iris.data[:, 0], iris.data[:, 1], c=iris.target)
    ax.set(xlabel=iris.feature_names[0], ylabel=iris.feature_names[1])
    _ = ax.legend(
        scatter.legend_elements()[0], iris.target_names, loc="lower right", title="Classes"
    )
    plt.show()

X, y = iris['data'], iris['target']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20, random_state=42)

model = LinearRegression()

model.fit(X_train, y_train)

score = model.score(X_test, y_test)

print(f"Score {score}")

output = model.predict(np.array([5.1, 3.5, 1.4, 0.2]).reshape(1, -1))     

# 0 : setosa
# 1 : versicolor
# 2 : virginica

if output < 0.5:
    predict = "setosa"
elif output < 1.5:
    predict = "versicolor"
else:
    predict = "virginica"

print(f"Predict : {predict} | Output : ", output)