Added saving predicted values to a .pred file

This commit is contained in:
Ilya Dziamidchyk 2024-01-12 01:00:07 +01:00
parent aa9388938d
commit 62d76361e7

View File

@ -8,20 +8,19 @@ TEST_DATA_DIR = "datasets_test"
test_df_list = [] test_df_list = []
for file in os.listdir(TEST_DATA_DIR): for file in os.listdir(TEST_DATA_DIR):
file_path = os.path.join(TEST_DATA_DIR, file) file_path = os.path.join(TEST_DATA_DIR, file)
df = pd.read_csv(file_path, delim_whitespace=True, skiprows=1, data_test = pd.read_csv(file_path, delim_whitespace=True, skiprows=1,
names=["tbid", "tphys", "r", "vr", "vt", "ik1", "ik2", "sm1", "sm2", "a", "e", names=["tbid", "tphys", "r", "vr", "vt", "ik1", "ik2", "sm1", "sm2", "a", "e",
"collapsed"]) "collapsed"])
test_df_list.append(df) X_test = data_test.iloc[:, 1:-1].values
data_test = pd.concat(test_df_list, ignore_index=True).sample(frac=1, random_state=42) imputer = SimpleImputer(strategy='mean')
X_test = data_test.iloc[:, 1:-1].values X_imputed = imputer.fit_transform(X_test)
imputer = SimpleImputer(strategy='mean') model_filename = 'trained_model.pkl'
X_imputed = imputer.fit_transform(X_test) model = joblib.load(model_filename)
model_filename = 'trained_model.pkl' predictions = model.predict(X_test)
model = joblib.load(model_filename) data_test['prediction'] = predictions
predictions = model.predict(X_test)
print(predictions)
save_path = os.path.join(TEST_DATA_DIR, file + '.pred')
data_test.to_csv(save_path, index=False, sep=' ')