Merge pull request 'Adding prediction function' (#18) from prediction into master

Reviewed-on: #18
pull/21/head
Dorian HODIN 10 months ago
commit b4bd477527

@ -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)
Loading…
Cancel
Save