Fix rmse saving
Some checks failed
s434804-evaluation/pipeline/head There was a failure building this commit

This commit is contained in:
Dawid 2021-05-14 21:25:04 +02:00
parent b0348f72f9
commit 81c0a9a7dd

View File

@ -12,17 +12,18 @@ alldata = pd.read_csv(
X = alldata['total_vaccinations'].to_numpy().ravel().reshape(-1, 1) X = alldata['total_vaccinations'].to_numpy().ravel().reshape(-1, 1)
y = alldata['daily_vaccinations'].to_numpy().ravel().reshape(-1, 1) y = alldata['daily_vaccinations'].to_numpy().ravel().reshape(-1, 1)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 6) X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2)
lin_reg = reg.fit(X_train, y_train) lin_reg = reg.fit(X_train, y_train)
score = lin_reg.score(X_test, y_test) score = lin_reg.score(X_test, y_test)
prediction = lin_reg.predict(X_test) prediction = lin_reg.predict(X_test)
build_no = sys.argv[1] if len(sys.argv) > 1 else 0 build_no = sys.argv[1] if len(sys.argv) > 1 else 0
data = {"rmse": [rmse], "build": [build_no]} rmse = rmse(y_test, prediction, squared=False)
data = {"rmse": [round(rmse, 2)], "build": [build_no]}
df = pd.DataFrame(data=data) df = pd.DataFrame(data=data)
with open("evaluation.csv", "a") as f: with open("evaluation.csv", "a") as f:
df.to_csv(f, header=f.tell() == 0, index=False) df.to_csv(f, header=f.tell() == 0, index=False)
print("RMSE:", rmse(y_test, prediction, squared=False)) print("RMSE:", rmse)
print("Score:", score) print("Score:", score)