{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"}],"dockerImageVersionId":30698,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Do Transformers Really Perform Bad for Graph Representation?\n# https://arxiv.org/abs/2106.05234\n# Tutorial: https://huggingface.co/blog/graphml-classification\n# I ran out of GPU quota so I changed BS to 2, original was 64\nfrom datasets import load_dataset\n\n# There is only one split on the hub\ndataset = load_dataset(\"OGB/ogbg-molhiv\")\n\ndataset = dataset.shuffle(seed=0)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-30T21:34:46.657254Z","iopub.execute_input":"2024-05-30T21:34:46.657795Z","iopub.status.idle":"2024-05-30T21:34:48.457079Z","shell.execute_reply.started":"2024-05-30T21:34:46.657749Z","shell.execute_reply":"2024-05-30T21:34:48.455781Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nnp.unique(np.array(dataset['train']['y']),return_counts=True)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:34:48.459550Z","iopub.execute_input":"2024-05-30T21:34:48.460050Z","iopub.status.idle":"2024-05-30T21:34:49.040137Z","shell.execute_reply.started":"2024-05-30T21:34:48.460008Z","shell.execute_reply":"2024-05-30T21:34:49.039055Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import networkx as nx\nimport matplotlib.pyplot as plt\n\n# We want to plot the first train graph\ngraph = dataset[\"train\"][0]\n\nedges = graph[\"edge_index\"]\nnum_edges = len(edges[0])\nnum_nodes = graph[\"num_nodes\"]\n\n# Conversion to networkx format\nG = nx.Graph()\nG.add_nodes_from(range(num_nodes))\nG.add_edges_from([(edges[0][i], edges[1][i]) for i in range(num_edges)])\n\n# Plot\nnx.draw(G)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:34:49.041350Z","iopub.execute_input":"2024-05-30T21:34:49.042324Z","iopub.status.idle":"2024-05-30T21:34:49.348374Z","shell.execute_reply.started":"2024-05-30T21:34:49.042290Z","shell.execute_reply":"2024-05-30T21:34:49.347020Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers.models.graphormer.collating_graphormer import preprocess_item, GraphormerDataCollator\n\ndataset_processed = dataset.map(preprocess_item, batched=False)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:34:49.351100Z","iopub.execute_input":"2024-05-30T21:34:49.351490Z","iopub.status.idle":"2024-05-30T21:34:49.523323Z","shell.execute_reply.started":"2024-05-30T21:34:49.351461Z","shell.execute_reply":"2024-05-30T21:34:49.521936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import GraphormerForGraphClassification\n\nmodel = GraphormerForGraphClassification.from_pretrained(\n    \"clefourrier/pcqm4mv2_graphormer_base\",\n    num_classes=2, # num_classes for the downstream task \n    ignore_mismatched_sizes=True,\n)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:34:49.525265Z","iopub.execute_input":"2024-05-30T21:34:49.526159Z","iopub.status.idle":"2024-05-30T21:34:53.274440Z","shell.execute_reply.started":"2024-05-30T21:34:49.526112Z","shell.execute_reply":"2024-05-30T21:34:53.272740Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import TrainingArguments, Trainer\n\ntraining_args = TrainingArguments(\n    \"graph-classification\",\n    logging_dir=\"graph-classification\",\n    per_device_train_batch_size=2,#64,\n    per_device_eval_batch_size=2,#64,\n    auto_find_batch_size=True, # batch size can be changed automatically to prevent OOMs\n    gradient_accumulation_steps=10,\n    dataloader_num_workers=4, #1, \n    num_train_epochs=20,\n    evaluation_strategy=\"epoch\",\n    logging_strategy=\"epoch\",\n    push_to_hub=False,\n)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:37:59.990763Z","iopub.execute_input":"2024-05-30T21:37:59.994861Z","iopub.status.idle":"2024-05-30T21:38:00.029224Z","shell.execute_reply.started":"2024-05-30T21:37:59.994612Z","shell.execute_reply":"2024-05-30T21:38:00.023618Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer = Trainer(\n    model=model,\n    args=training_args,\n    train_dataset=dataset_processed[\"train\"],\n    eval_dataset=dataset_processed[\"validation\"],\n    data_collator=GraphormerDataCollator(),\n)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:38:02.615719Z","iopub.execute_input":"2024-05-30T21:38:02.617121Z","iopub.status.idle":"2024-05-30T21:38:02.711041Z","shell.execute_reply.started":"2024-05-30T21:38:02.616998Z","shell.execute_reply":"2024-05-30T21:38:02.706500Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for b in trainer.get_train_dataloader():\n    print(b)\n    break","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:40:24.668915Z","iopub.execute_input":"2024-05-30T21:40:24.669465Z","iopub.status.idle":"2024-05-30T21:40:25.609174Z","shell.execute_reply.started":"2024-05-30T21:40:24.669428Z","shell.execute_reply":"2024-05-30T21:40:25.607750Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"b.keys()","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:40:33.315304Z","iopub.execute_input":"2024-05-30T21:40:33.315772Z","iopub.status.idle":"2024-05-30T21:40:33.326351Z","shell.execute_reply.started":"2024-05-30T21:40:33.315734Z","shell.execute_reply":"2024-05-30T21:40:33.324762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output = model(\n    b['input_nodes'],\n    b['input_edges'],\n    b['attn_bias'],\n    b['in_degree'],\n    b['out_degree'],\n    b['spatial_pos'],\n    b['attn_edge_type']\n)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:42:11.149914Z","iopub.execute_input":"2024-05-30T21:42:11.150562Z","iopub.status.idle":"2024-05-30T21:42:11.356291Z","shell.execute_reply.started":"2024-05-30T21:42:11.150514Z","shell.execute_reply":"2024-05-30T21:42:11.354555Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output['logits']","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:42:14.837095Z","iopub.execute_input":"2024-05-30T21:42:14.837622Z","iopub.status.idle":"2024-05-30T21:42:14.851926Z","shell.execute_reply.started":"2024-05-30T21:42:14.837584Z","shell.execute_reply":"2024-05-30T21:42:14.850344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.train()","metadata":{"execution":{"iopub.status.busy":"2024-05-30T21:45:45.339727Z","iopub.execute_input":"2024-05-30T21:45:45.340269Z","iopub.status.idle":"2024-05-30T21:57:22.599157Z","shell.execute_reply.started":"2024-05-30T21:45:45.340231Z","shell.execute_reply":"2024-05-30T21:57:22.596612Z"},"trusted":true},"execution_count":null,"outputs":[]}]}