From 2cbcbeaf7ecb10808a3ee97f42ad4bdf019b047e Mon Sep 17 00:00:00 2001 From: dorian Date: Fri, 21 Jun 2024 14:31:19 +0200 Subject: [PATCH] Adding prediction file --- src/back/prediction.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 src/back/prediction.py diff --git a/src/back/prediction.py b/src/back/prediction.py new file mode 100644 index 0000000..09c7556 --- /dev/null +++ b/src/back/prediction.py @@ -0,0 +1,19 @@ +from sklearn.model_selection import train_test_split +from sklearn.linear_model import LinearRegression +from sklearn.ensemble import RandomForestRegressor + + +def getColumnsForPredictionAndPredict(df,columns, columnGoal, algoOfPrediction): + predictors = df[columns] + target = df[columnGoal] + X_train, X_test, y_train, y_test = train_test_split(predictors, target, test_size=0.2, random_state=42) + + if algoOfPrediction == "Régression Linéaire": + model = LinearRegression() + elif algoOfPrediction == "Forêt Aléatoire": + model = RandomForestRegressor(n_estimators=100) + else: + raise NameError("No method name : \"" + algoOfPrediction + "\"") + + model.fit(X_train, y_train) + return model.predict(X_test) \ No newline at end of file