<jupyter_start><jupyter_text>Knowledge graph link prediction with ComplEx Run the latest release of this notebook: This notebook reproduces the experiments done in the paper that introduced the ComplEx algorithm: Complex Embeddings for Simple Link Prediction, <NAME>, <NAME>, <NAME>, <NAME> and <NAME>, ICML 2016. http://jmlr.org/proceedings/papers/v48/trouillon16.pdfIn table 2, the paper reports five metrics measured on the WN18 and FB15K datasets: "raw" MRR (mean reciprocal rank), "filtered" MRR and filtered Hits at {1, 3, 10}. This notebook measures all of these, as well as raw Hits at {1, 3, 10}.<jupyter_code># install StellarGraph if running on Google Colab
import sys
if 'google.colab' in sys.modules:
  %pip install -q stellargraph[demos]==1.3.0b
# verify that we're using the correct version of StellarGraph for this notebook
import stellargraph as sg

try:
    sg.utils.validate_notebook_version("1.3.0b")
except AttributeError:
    raise ValueError(
        f"This notebook requires StellarGraph version 1.3.0b, but a different version {sg.__version__} is installed.  Please see <https://github.com/stellargraph/stellargraph/issues/1172>."
    ) from None
from stellargraph import datasets, utils
from tensorflow.keras import callbacks, optimizers, losses, metrics, regularizers, Model
import numpy as np
import pandas as pd

from stellargraph.mapper import KGTripleGenerator
from stellargraph.layer import ComplEx

from IPython.display import HTML<jupyter_output><empty_output><jupyter_text>InitialisationWe need to set up our model parameters, like the number of epochs to train for, and the dimension of the embedding vectors we compute for each node and for each edge type.The evaluation is performed in three steps:1. Load the data2. Train a model3. Evaluate the modelThe paper says that it used:- the AdaGrad optimiser for 1000 epochs with an early stopping criterion evaluated every 50 epochs, but we've found using the Adam optimiser allows for much fewer epochs- an embedding dimension of 150 or 200, since they had close results- 10 negative samples (corrupted edges) per positive edge, which gives noticeably improved performance on FB15k compared to using 1, and but not for WN18 (the paper evaluated 1, 2, 5 and 10 negative samples)<jupyter_code>epochs = 50
embedding_dimension = 200
negative_samples = 10<jupyter_output><empty_output><jupyter_text>WN18The paper uses the WN18 and FB15k datasets for validation. These datasets are not good for evaluating algorithms because they contain "inverse relations", where `(s, r1, o)` implies `(o, r2, s)` for a pair of relation types `r1` and `r2` (for instance, `_hyponym` ("is more specific than") and `_hypernym` ("is more general than") in WN18), however, they work fine to demonstrate StellarGraph's functionality, and are appropriate to compare against the published results. Load the dataThe dataset comes with a defined train, test and validation split, each consisting of subject, relation, object triples. We can load a `StellarGraph` object with all of the triples, as well as the individual splits as Pandas DataFrames, using the `load` method of the `WN18` dataset. (See [the "Loading from Pandas" demo](../basics/loading-pandas.ipynb) for details on how data can be loaded.)<jupyter_code>wn18 = datasets.WN18()
display(HTML(wn18.description))
wn18_graph, wn18_train, wn18_test, wn18_valid = wn18.load()
print(wn18_graph.info())<jupyter_output>StellarDiGraph: Directed multigraph
 Nodes: 40943, Edges: 151442

 Node types:
  default: [40943]
    Features: none
    Edge types: default-_also_see->default, default-_derivationally_related_form->default, default-_has_part->default, default-_hypernym->default, default-_hyponym->default, ... (13 more)

 Edge types:
    default-_hyponym->default: [37221]
    default-_hypernym->default: [37221]
    default-_derivationally_related_form->default: [31867]
    default-_member_meronym->default: [7928]
    default-_member_holonym->default: [7928]
    default-_part_of->default: [5148]
    default-_has_part->default: [5142]
    default-_member_of_domain_topic->default: [3341]
    default-_synset_domain_topic_of->default: [3335]
    default-_instance_hyponym->default: [3150]
    default-_instance_hypernym->default: [3150]
    default-_also_see->default: [1396]
    default-_verb_group->default: [1220]
    default-_member_of_domain_region->default: [983]
    default-_synset_domain_region_of->defa[...]<jupyter_text>Train a modelThe ComplEx algorithm consists of some embedding layers and a scoring layer, but the `ComplEx` object means these details are invisible to us. The `ComplEx` model consumes "knowledge-graph triples", which can be produced in the appropriate format using `KGTripleGenerator`.<jupyter_code>wn18_gen = KGTripleGenerator(
    wn18_graph, batch_size=len(wn18_train) // 100  # ~100 batches per epoch
)

wn18_complex = ComplEx(
    wn18_gen,
    embedding_dimension=embedding_dimension,
    embeddings_regularizer=regularizers.l2(1e-7),
)

wn18_inp, wn18_out = wn18_complex.in_out_tensors()

wn18_model = Model(inputs=wn18_inp, outputs=wn18_out)

wn18_model.compile(
    optimizer=optimizers.Adam(lr=0.001),
    loss=losses.BinaryCrossentropy(from_logits=True),
    metrics=[metrics.BinaryAccuracy(threshold=0.0)],
)<jupyter_output><empty_output><jupyter_text>Inputs for training are produced by calling the `KGTripleGenerator.flow` method, this takes a dataframe with `source`, `label` and `target` columns, where each row is a true edge in the knowledge graph.  The `negative_samples` parameter controls how many random edges are created for each positive edge to use as negative examples for training.<jupyter_code>wn18_train_gen = wn18_gen.flow(
    wn18_train, negative_samples=negative_samples, shuffle=True
)
wn18_valid_gen = wn18_gen.flow(wn18_valid, negative_samples=negative_samples)
wn18_es = callbacks.EarlyStopping(monitor="val_loss", patience=10)
wn18_history = wn18_model.fit(
    wn18_train_gen, validation_data=wn18_valid_gen, epochs=epochs, callbacks=[wn18_es]
)
utils.plot_history(wn18_history)<jupyter_output><empty_output><jupyter_text>Evaluate the modelWe've now trained a model, so we can apply the evaluation procedure from the paper to it. This is done by taking each test edge `E = (s, r, o)`, and scoring it against all mutations `(s, r, n)` and `(n, r, o)` for every node `n` in the graph, that is, doing a prediction for every one of these edges similar to `E`. The "raw" rank is the number of mutated edges that have a higher predicted score than the true `E`.<jupyter_code>wn18_raw_ranks, wn18_filtered_ranks = wn18_complex.rank_edges_against_all_nodes(
    wn18_gen.flow(wn18_test), wn18_graph
)
# helper function to compute metrics from a dictionary of name -> array of ranks
def results_as_dataframe(name_to_results):
    return pd.DataFrame(
        name_to_results.values(),
        columns=["mrr", "hits at 1", "hits at 3", "hits at 10"],
        index=name_to_results.keys(),
    )


def summarise(name_to_ranks):
    return results_as_dataframe(
        {
            name: (
                np.mean(1 / ranks),
                np.mean(ranks <= 1),
                np.mean(ranks < 3),
                np.mean(ranks <= 10),
            )
            for name, ranks in name_to_ranks.items()
        }
    )
summarise({"raw": wn18_raw_ranks, "filtered": wn18_filtered_ranks})<jupyter_output><empty_output><jupyter_text>For comparison, Table 2 in the paper gives the following results for WN18 (`NaN` denotes values the paper does not include). All of the numbers are similar:<jupyter_code>results_as_dataframe(
    {"raw": (0.587, None, None, None), "filtered": (0.941, 0.936, 0.945, 0.947)}
)<jupyter_output><empty_output><jupyter_text>FB15kNow that we know the process, we can apply the model on the FB15k dataset in the same way. Loading the data<jupyter_code>fb15k = datasets.FB15k()
display(HTML(fb15k.description))
fb15k_graph, fb15k_train, fb15k_test, fb15k_valid = fb15k.load()
print(fb15k_graph.info())<jupyter_output>StellarDiGraph: Directed multigraph
 Nodes: 14951, Edges: 592213

 Node types:
  default: [14951]
    Features: none
    Edge types: default-/american_football/football_coach/coaching_history./american_football/football_historical_coach_position/position->default, default-/american_football/football_coach/coaching_history./american_football/football_historical_coach_position/team->default, default-/american_football/football_coach_position/coaches_holding_this_position./american_football/football_historical_coach_position/coach->default, default-/american_football/football_coach_position/coaches_holding_this_position./american_football/football_historical_coach_position/team->default, default-/american_football/football_player/current_team./american_football/football_roster_position/position->default, ... (1340 more)

 Edge types:
    default-/award/award_nominee/award_nominations./award/award_nomination/award_nominee->default: [19764]
    default-/film/film/release_date_s./film/film_r[...]<jupyter_text>Train a model<jupyter_code>fb15k_gen = KGTripleGenerator(
    fb15k_graph, batch_size=len(fb15k_train) // 100  # ~100 batches per epoch
)

fb15k_complex = ComplEx(
    fb15k_gen,
    embedding_dimension=embedding_dimension,
    embeddings_regularizer=regularizers.l2(1e-8),
)

fb15k_inp, fb15k_out = fb15k_complex.in_out_tensors()

fb15k_model = Model(inputs=fb15k_inp, outputs=fb15k_out)
fb15k_model.compile(
    optimizer=optimizers.Adam(lr=0.001),
    loss=losses.BinaryCrossentropy(from_logits=True),
    metrics=[metrics.BinaryAccuracy(threshold=0.0)],
)
fb15k_train_gen = fb15k_gen.flow(
    fb15k_train, negative_samples=negative_samples, shuffle=True
)
fb15k_valid_gen = fb15k_gen.flow(fb15k_valid, negative_samples=negative_samples)
fb15k_es = callbacks.EarlyStopping(monitor="val_loss", patience=10)
fb15k_history = fb15k_model.fit(
    fb15k_train_gen, validation_data=fb15k_valid_gen, epochs=epochs, callbacks=[fb15k_es]
)
utils.plot_history(fb15k_history)<jupyter_output><empty_output><jupyter_text>Evaluate the model<jupyter_code>fb15k_raw_ranks, fb15k_filtered_ranks = fb15k_complex.rank_edges_against_all_nodes(
    fb15k_gen.flow(fb15k_test), fb15k_graph
)
summarise({"raw": fb15k_raw_ranks, "filtered": fb15k_filtered_ranks})<jupyter_output><empty_output><jupyter_text>For comparison, Table 2 in the paper gives the following results for FB15k:<jupyter_code>results_as_dataframe(
    {"raw": (0.242, None, None, None), "filtered": (0.692, 0.599, 0.759, 0.850)}
)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09302105063531196, 'fraction_numerical': 0.04210127062393324, 'mean_word_length': 5.526608910891089, 'pattern_counts': {'":': 8, '<': 52, '<?xml version=': 0, '>': 76, 'https://': 1, '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}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '97048', 'license': 'Apache-2.0', 'n_tokens_mistral': 3780, 'n_tokens_neox': 3411, 'n_words': 937, 'path': 'demos/link-prediction/complex-link-prediction.ipynb', 'repo_name': 'vishalbelsare/stellargraph', 'size': '10591'}