FlagRelease/materials.smi-ted-nvidia-FlagOS

1import sys2sys.path.append("../models")3sys.path.append("../")4 5import models.fm4m as fm4m6import pandas as pd7import numpy as np8from sklearn.svm import SVR9from sklearn.compose import TransformedTargetRegressor10from sklearn.preprocessing import MinMaxScaler11from sklearn.metrics import mean_squared_error12 13train_df  = pd.read_csv(f"../data/lce/train.csv").dropna()14test_df  = pd.read_csv(f"../data/lce/test.csv").dropna()15 16 17# Make a list of smiles18train_smiles_list = pd.concat([train_df[f'smi{i}'] for i in range(1, 7)]).unique().tolist()19test_smiles_list = pd.concat([test_df[f'smi{i}'] for i in range(1, 7)]).unique().tolist()20 21fm4m.avail_models()22 23model_type = "SMI-TED"24train_emb, test_emb = fm4m.get_representation(train_smiles_list,test_smiles_list, model_type, return_tensor=False)25 26train_emb = [np.nan if row.isna().all() else row.dropna().tolist() for _, row in train_emb.iterrows()]27test_emb = [np.nan if row.isna().all() else row.dropna().tolist() for _, row in test_emb.iterrows()]28 29train_dict = dict(zip(train_smiles_list, train_emb))30test_dict = dict(zip(test_smiles_list, test_emb))31 32def replace_with_list(value, my_dict):33    return my_dict.get(value, value)34 35# Replacement the smiles string with its embeddings36df_train_emb = train_df.applymap(lambda x: replace_with_list(x, train_dict))37df_test_emb = test_df.applymap(lambda x: replace_with_list(x, test_dict))38 39# Drop rows with NaN and reset index40df_train_emb = df_train_emb.dropna().reset_index(drop=True)41df_test_emb = df_test_emb.dropna().reset_index(drop=True)42 43# Define a function to handle repetitive tasks44def compute_components(df, smi_cols, conc_cols):45    components = [df[smi].apply(pd.Series).mul(df[conc], axis=0) for smi, conc in zip(smi_cols, conc_cols)]46    return sum(components)47 48# List of columns to process49smi_cols = [f'smi{i}' for i in range(1, 7)]50conc_cols = [f'conc{i}' for i in range(1, 7)]51 52# Train data processing53x_train = compute_components(df_train_emb, smi_cols, conc_cols)54y_train = pd.DataFrame(df_train_emb["LCE"], columns=["LCE"])55 56# Test data processing57X_test = compute_components(df_test_emb, smi_cols, conc_cols)58y_test = pd.DataFrame(df_test_emb["LCE"], columns=["LCE"])59 60regressor = SVR(kernel="rbf", degree=3, C=5, gamma="scale", epsilon=0.01)61model = TransformedTargetRegressor(regressor=regressor,62                                   transformer=MinMaxScaler(feature_range=(-1, 1))63                                   ).fit(x_train, y_train)64 65y_prob = model.predict(X_test)66RMSE_score = mean_squared_error(y_test, y_prob, squared=False)67print(RMSE_score)68 69