<commit_before>"""
Link Attribute Inference on a Heterogeneous Graph for the MovieLens dataset

usage: movielens.py [-h] [-c [CHECKPOINT]] [-n BATCH_SIZE] [-e EPOCHS]
                    [-s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]]
                    [-l [LAYER_SIZE [LAYER_SIZE ...]]] [-m METHOD] [-b]
                    [-g GRAPH] [-f FEATURES] [-t TARGET]

Run GraphSAGE on movielens

optional arguments:
  -h, --help            show this help message and exit
  -c [CHECKPOINT], --checkpoint [CHECKPOINT]
                        Load a save checkpoint file
  -n BATCH_SIZE, --batch_size BATCH_SIZE
                        Load a save checkpoint file
  -e EPOCHS, --epochs EPOCHS
                        Number of epochs to train for
  -s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]], --neighbour_samples [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]
                        The number of nodes sampled at each layer
  -l [LAYER_SIZE [LAYER_SIZE ...]], --layer_size [LAYER_SIZE [LAYER_SIZE ...]]
                        The number of hidden features at each layer
  -m METHOD, --method METHOD
                        The edge regression method: 'concat', 'mul', or 'ip
  -b, --baseline        Use a learned offset for each node.
  -g GRAPH, --graph GRAPH
                        The graph stored in networkx pickle format.
  -f FEATURES, --features FEATURES
                        The node features to use, stored as a pickled numpy
                        array.
  -t TARGET, --target TARGET
                        The target edge attribute, default is 'score'
"""
import time
import os
import argparse
import pickle
import random

import tensorflow as tf
import numpy as np
import pandas as pd
import networkx as nx
import itertools as it
import functools as ft

from collections import defaultdict
from typing import AnyStr, Any, List, Tuple, Dict, Optional

from keras import backend as K
from keras import Input, Model, optimizers, losses, activations, metrics
from keras.layers import Layer, Dense, Concatenate, Multiply, Activation, Lambda
from keras.utils import Sequence

from recommender import MovielensData, root_mean_square_error, regression_predictor


def train(ml_data: MovielensData, batch_size: int = 1000, num_epochs: int = 10):
    # Create data iterators
    train_iter, test_iter = ml_data.edge_data_generators(batch_size)

    # HINSAGE model
    hs = ml_data.create_model()

    x_inp = [Input(shape=s) for s in ml_data.shapes_from_schema()]
    x_out = hs(x_inp)

    # Final estimator layer
    score_prediction = regression_predictor(
        hidden_1=32, hidden_2=32, method=ml_data.edge_regressor
    )(x_out)

    # Create Keras model for training
    model = Model(inputs=x_inp, outputs=score_prediction)
    model.compile(
        optimizer=optimizers.Adam(lr=0.0005),
        loss=losses.mean_squared_error,
        metrics=[root_mean_square_error, metrics.mae],
    )

    # Train model
    history = model.fit_generator(
        train_iter, epochs=num_epochs, verbose=2, shuffle=True
    )

    # Evaluate and print metrics
    test_metrics = model.evaluate_generator(test_iter)

    print("Test Evaluation:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))


def test(ml_data: MovielensData, model_file: AnyStr):
    pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run GraphSAGE on movielens")
    parser.add_argument(
        "-c",<commit_msg>Original movielens example: changed the link regression to not apply dense transforms to edge features, for easy comparison with the new movielens example<commit_after>"""
Link Attribute Inference on a Heterogeneous Graph for the MovieLens dataset

usage: movielens.py [-h] [-c [CHECKPOINT]] [-n BATCH_SIZE] [-e EPOCHS]
                    [-s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]]
                    [-l [LAYER_SIZE [LAYER_SIZE ...]]] [-m METHOD] [-b]
                    [-g GRAPH] [-f FEATURES] [-t TARGET]

Run GraphSAGE on movielens

optional arguments:
  -h, --help            show this help message and exit
  -c [CHECKPOINT], --checkpoint [CHECKPOINT]
                        Load a save checkpoint file
  -n BATCH_SIZE, --batch_size BATCH_SIZE
                        Load a save checkpoint file
  -e EPOCHS, --epochs EPOCHS
                        Number of epochs to train for
  -s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]], --neighbour_samples [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]
                        The number of nodes sampled at each layer
  -l [LAYER_SIZE [LAYER_SIZE ...]], --layer_size [LAYER_SIZE [LAYER_SIZE ...]]
                        The number of hidden features at each layer
  -m METHOD, --method METHOD
                        The edge regression method: 'concat', 'mul', or 'ip
  -b, --baseline        Use a learned offset for each node.
  -g GRAPH, --graph GRAPH
                        The graph stored in networkx pickle format.
  -f FEATURES, --features FEATURES
                        The node features to use, stored as a pickled numpy
                        array.
  -t TARGET, --target TARGET
                        The target edge attribute, default is 'score'
"""
import time
import os
import argparse
import pickle
import random

import tensorflow as tf
import numpy as np
import pandas as pd
import networkx as nx
import itertools as it
import functools as ft

from collections import defaultdict
from typing import AnyStr, Any, List, Tuple, Dict, Optional

from keras import backend as K
from keras import Input, Model, optimizers, losses, activations, metrics
from keras.layers import Layer, Dense, Concatenate, Multiply, Activation, Lambda
from keras.utils import Sequence

from recommender import MovielensData, root_mean_square_error, regression_predictor


def train(ml_data: MovielensData, batch_size: int = 1000, num_epochs: int = 10):
    # Create data iterators
    train_iter, test_iter = ml_data.edge_data_generators(batch_size)

    # HINSAGE model
    hs = ml_data.create_model()

    x_inp = [Input(shape=s) for s in ml_data.shapes_from_schema()]
    x_out = hs(x_inp)

    # Final estimator layer
    score_prediction = regression_predictor(
        hidden_1=None, hidden_2=None, method=ml_data.edge_regressor
    )(x_out)

    # Create Keras model for training
    model = Model(inputs=x_inp, outputs=score_prediction)
    model.compile(
        optimizer=optimizers.Adam(lr=0.0005),
        loss=losses.mean_squared_error,
        metrics=[root_mean_square_error, metrics.mae],
    )

    # Train model
    history = model.fit_generator(
        train_iter, epochs=num_epochs, verbose=2, shuffle=True
    )

    # Evaluate and print metrics
    test_metrics = model.evaluate_generator(test_iter)

    print("Test Evaluation:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))


def test(ml_data: MovielensData, model_file: AnyStr):
    pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run GraphSAGE on movielens")
    parser.add_argument(
        "-c",
{'fraction_non_alphanumeric': 0.08711535759449206, 'fraction_numerical': 0.005058311086131797, 'mean_word_length': 2.6823590274185203, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3513792', 'n_tokens_mistral': 2239, 'n_tokens_neox': 2118, 'n_words': 691}