<jupyter_start><jupyter_text>1. run all models - extract val and test set2. do mean of val and test set<jupyter_code>import pandas as pd
from ogb.lsc import PygPCQM4MDataset, PCQM4MEvaluator
import numpy as np

from torch_geometric.data import DataLoader
from tqdm import tqdm
import torch

import os
"""
LOAD HYPERPARAMETERS
"""
# get args from main_gnn CLI
class Argument(object):
    name = "args"
    
args = Argument()
args.batch_size = 256
args.num_workers = 0
args.num_layers = 5
args.emb_dim = 600
args.drop_ratio = 0
args.graph_pooling = "sum"
args.device = 0

shared_params = {
    'num_layers': args.num_layers,
    'emb_dim': args.emb_dim,
    'drop_ratio': args.drop_ratio,
    'graph_pooling': args.graph_pooling
}

device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")<jupyter_output><empty_output><jupyter_text>load data<jupyter_code>"""
LOAD DATA
"""
### importing OGB-LSC
from ogb.lsc import PygPCQM4MDataset, PCQM4MEvaluator

dataset = PygPCQM4MDataset(root = 'dataset/')
split_idx = dataset.get_idx_split()

valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=args.batch_size, shuffle=False, num_workers = args.num_workers)
test_loader = DataLoader(dataset[split_idx["test"]], batch_size=args.batch_size, shuffle=False, num_workers = args.num_workers)
print(len(split_idx["train"]), len(split_idx["test"]), len(split_idx["valid"]))
def get_prediction(model, loader):
    y_true = []
    y_pred = []
    for step, batch in enumerate(tqdm(loader, desc="Iteration")):

        # put batch to cuda
        batch = batch.to(device)

        pred = model(batch)
        pred = pred.view(-1)    
        # collate prediction
        pred_np = pred.detach().cpu().numpy()
        y_pred.extend(pred_np)

        # collate label
        label = batch.y.detach().cpu().numpy()
        y_true.extend(label)
        
    return y_true, y_pred<jupyter_output><empty_output><jupyter_text>pygeom models<jupyter_code>"""
gin-virtual-bayes-lastLayer
"""
from gnn import BayesianGNN

args.checkpoint_dir = "models/gin-virtual-bnn-lastLayer/checkpoint"

"""
LOAD Checkpoint data
"""
checkpoint = torch.load(os.path.join(args.checkpoint_dir, 'checkpoint.pt'))

gnn_name = "gin-virtual-bnn-lastLayer"
gnn_type = "gin"
virtual_node = True

model = BayesianGNN(gnn_type = gnn_type, virtual_node = virtual_node, last_layer_only=True, **shared_params).to(device)
model.load_state_dict(checkpoint["model_state_dict"], strict=True)
model.state_dict()
# set model to evaluation-only
model.eval()
print(gnn_name, checkpoint["best_val_mae"])

"""
get prediction for val
"""
y_true, y_pred = get_prediction(model, valid_loader)    
evaluator = PCQM4MEvaluator()
input_dict = {"y_true": torch.tensor(y_true), "y_pred": torch.tensor(y_pred)}
print("BayesianGNN", evaluator.eval(input_dict)["mae"])

# save results to dataframe
ginBNNDF = pd.DataFrame(zip(split_idx["valid"].cpu().tolist(), y_true, y_pred), columns=["molecule_idx", "label", "gin-pred"])

"""
get prediction for test
"""
_, ginBNN_test_pred = get_prediction(model, test_loader)    
ginBNN_test_pred = np.array(ginBNN_test_pred)
ginBNN_test_pred.shape
"""
gin-virtual
"""
from gnn import GNN

args.checkpoint_dir = "models/gin-virtual/checkpoint"

"""
LOAD Checkpoint data
"""
checkpoint = torch.load(os.path.join(args.checkpoint_dir, 'checkpoint.pt'))

gnn_name = "gin-virtual"
gnn_type = "gin"
virtual_node = True

model = GNN(gnn_type = gnn_type, virtual_node = virtual_node, **shared_params).to(device)
model.load_state_dict(checkpoint["model_state_dict"], strict=True)
model.state_dict()
# set model to evaluation-only
model.eval()
print(gnn_name, checkpoint["best_val_mae"])

"""
get prediction for val
"""
y_true, y_pred = get_prediction(model, valid_loader)    
evaluator = PCQM4MEvaluator()
input_dict = {"y_true": torch.tensor(y_true), "y_pred": torch.tensor(y_pred)}
print("BayesianGNN", evaluator.eval(input_dict)["mae"])

# save results to dataframe
ginDF = pd.DataFrame(zip(split_idx["valid"].cpu().tolist(), y_true, y_pred), columns=["molecule_idx", "label", "gin-pred"])

"""
get prediction for test
"""
_, gin_test_pred = get_prediction(model, test_loader)  
gin_test_pred = np.array(gin_test_pred)
gin_test_pred.shape<jupyter_output>Iteration:   0%|                                                                              | 0/1487 [00:00<?, ?it/s]<jupyter_text>load dgl models<jupyter_code># get valid result
ginDiffPoolDF = pd.read_csv("../pcqm4m-dgl/models/gin-virtual-diffpool/gin-virtual-diffPool-validResult.csv")

# get test result
with open('../pcqm4m-dgl/models/gin-virtual-diffpool/gin-virtual-diffpool-testResult.npy', 'rb') as f:
    ginDiffPool_test_pred = np.load(f)<jupyter_output><empty_output><jupyter_text>ensembling<jupyter_code>ginDF.head()
ginBNNDF.head()
ginDiffPoolDF.head()
"""
valid ensembling
"""
for name, pairs in [
    ["gin", (ginDF["gin-pred"], ginDF["gin-pred"])],
    ["gin-bnn", (ginBNNDF["gin-pred"], ginBNNDF["gin-pred"])],
    ["gin-diffpool", (ginDiffPoolDF["gin-diffpool-pred"], ginDiffPoolDF["gin-diffpool-pred"])],
    ["gin + gin-bnn", (ginDF["gin-pred"], ginBNNDF["gin-pred"])],
    ["gin + gin-diffpool", (ginDF["gin-pred"], ginDiffPoolDF["gin-diffpool-pred"])],
    ["gin-bnn + gin-diffpool", (ginBNNDF["gin-pred"], ginDiffPoolDF["gin-diffpool-pred"])],
    ["gin + gin-bnn + gin-diffpool", (ginDF["gin-pred"], ginBNNDF["gin-pred"], ginDiffPoolDF["gin-diffpool-pred"])],
]:

    mean_pred = np.mean(list(zip(*pairs)), axis=1)
    evaluator = PCQM4MEvaluator()
    input_dict = {"y_true": torch.tensor(y_true), "y_pred": torch.tensor(mean_pred)}
    print(name, evaluator.eval(input_dict)["mae"])

"""
test ensembling - and save
"""

args.save_test_dir = "models/ensembling-gin-virtual-diffpool-bayesLastLayer"

pairs = (ginBNN_test_pred, gin_test_pred, ginDiffPool_test_pred)
ensemble_pred = np.mean(list(zip(*pairs)), axis=1)

evaluator.save_test_submission({'y_pred': ensemble_pred}, args.save_test_dir)

len(ensemble_pred)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15286518714806227, 'fraction_numerical': 0.005796621397813845, 'mean_word_length': 5.245087900723888, 'pattern_counts': {'":': 6, '<': 21, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '619940', 'license': 'MIT', 'n_tokens_mistral': 2271, 'n_tokens_neox': 2218, 'n_words': 380, 'path': 'examples/lsc/pcqm4m/ensemble-pcqm4m-pygeom.ipynb', 'repo_name': 'edwardelson/ogb', 'size': '6038'}