artcraft/models/tts/NeMo/tutorials/nlp/Neural_Machine_Translation-Text2Sparql.ipynb

In [ ]:

"""
You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.

Instructions for setting up Colab are as follows:
1. Open a new Python 3 notebook.
2. Import this notebook from GitHub (File -> Upload Notebook -> "GITHUB" tab -> copy/paste GitHub URL)
3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select "GPU" for hardware accelerator)
4. Run this cell to set up dependencies.
"""
# If you're using Google Colab and not running locally, run this cell

# install NeMo
BRANCH = 'v1.0.2'
!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]

In [ ]:

# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:
# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'

! pip install ipywidgets
! jupyter nbextension enable --py widgetsnbextension

# Please restart the kernel after running this cell

In [ ]:

import os
import wget 
import torch
import pytorch_lightning as pl
from omegaconf import OmegaConf

from nemo.collections import nlp as nemo_nlp
from nemo.utils.exp_manager import exp_manager
from nemo.collections.nlp.models import NeuralMachineTranslationModel

Motivation

RDF - based knowledge graphs (KG) are dense and flexible structures that allow rich representations of real data. Due to their complex nature, searching through their content can be difficult and requires an in-depth understanding of query syntax and graph specific properties. This problem can be solved via a natural language to query language translation model, which allows widespread use of even the most complicated KG.

Task Description

Neural Machine Translation (NMT) is the task of converting a sequence to another sequence using a neural network. While its most common use case is for translating from one language to another, we can also apply it to the task of converting a natural language request into a RDF query.

Dataset

To train our model, we use Text2Sparql, a synthetic dataset of natural language and WikiData SPARQL queries pairs generated with very little human intervention.

Some sentence-query examples from Text2Sparql:

Who is the mother of the director of Pulp Fiction?
SELECT ?end WHERE { [ Pulp Fiction ] wdt:P5 / wdt:P25 ?end . }

Is John Steinbeck the author of Green Eggs and Ham?
ASK { BIND ( [ John Steinbeck ] as ?end ) . [ Green Eggs and Ham ] wdt:P50 ?end . }

How many awards does the producer of Fast and Furious have?
SELECT ( COUNT ( DISTINCT ?end ) as ?endcount ) WHERE { [ Fast and Furious ] wdt:P162 / wdt:P166 ?end . }

Each query label can be broken down into several components:

• SPARQL syntax typically denoted in capital letters (eg. SELECT, ASK, BIND, COUNT, DISTINCT, WHERE)
• Variables, prefixed by a question mark. (eg: ?end)
• Property values, prefixed by P. Property values map the relationship between entities and are learned by the neural network. (eg: wdt:P5)
• Item values, prefixed by Q (hence we refer to them as q-values). Q-values define an entity and are not learned since fitting all possible items into the training set is impractical. Instead, they are represented as strings within brackets. (eg: [ Pulp Fiction ])

It is important to note that a query cannot be executed until its item representation is replaced with it's appropriate q-value. This step will be handled in the final section.

Download and Preprocess Data

The Text2Sparql dataset contains 3 files:

• train_queries_v3.tsv
• test_easy_queries_v3.tsv
• test_hard_queries_v3.tsv

with the format:

• [english] [tab] [sparql] [tab] [unique hash]

To download and convert the datasets for NeuralMachineTranslationDataset, we will reduce the table to 2 columns:

• [sentence] [tab] [label]

In [ ]:

# set the following paths
DATA_DIR = "PATH_TO_DATA"
WORK_DIR = "PATH_TO_CHECKPOINTS_AND_LOGS"

# NeMo Version
BRANCH = 'v1.0.2'

In [ ]:

# download import_datasets.py script to download and preprocess Text2Sparql
os.makedirs(WORK_DIR, exist_ok=True)
if not os.path.exists(WORK_DIR + "/get_squad.py"):
    print("Downloading import_datasets.py...")
    wget.download(f"https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/text2sparql/data/import_datasets.py", WORK_DIR)
else:
    print("import_datasets.py already exists")

In [ ]:

# run script
! python $WORK_DIR/import_datasets.py --source_data_dir $DATA_DIR --target_data_dir $DATA_DIR

Data and Model Parameters

In the following, we need to adjust the default model configuration for the NeMo experiment.

Note: This is just a baseline model for Text2Sparql to showcase usage and is not optimized for accuracy.

In [ ]:

# This is the model configuration file that we will download, do not change this
MODEL_CONFIG = "text2sparql_config.yaml"

# training parameters
GPUS = 1 if torch.cuda.is_available() else 0 # 0 for CPU, or list of GPU indicies
MAX_EPOCHS = 2 # number of training epochs

# model parameters, play with these
NEMO_PATH = f"{WORK_DIR}/bart.nemo"  # specify model savepath (models are saved as .nemo files)
BATCH_SIZE = 16
MAX_SEQ_LENGTH = 150

# specify seq2seq model you want to use
PRETRAINED_BART_MODEL = "facebook/bart-base"
TOKENIZER_NAME = "facebook/bart-base"

TRAIN_FILE = f"{DATA_DIR}/train.tsv"
EVAL_FILE = f"{DATA_DIR}/test_easy.tsv"
TEST_FILE = f"{DATA_DIR}/test_easy.tsv"  # can change to test_hard.tsv

# training parameters
LEARNING_RATE = 0.00004

Model Configuration

A NeuralMachineTranslationModel's config file declares multiple import sections. They are:

• trainer: Arguments to be passed to PyTorch Lightning
• model: All arguments that relate to the Model - language_model, tokenizers, datasets, optimizer, generate
• exp_manager: Arguments to be passed to NeMo's experiment manager
• hydra: Arguments to be passed to Hydra

In [ ]:

# download the model's default configuration file 
config_dir = WORK_DIR + "/configs/"
os.makedirs(config_dir, exist_ok=True)
if not os.path.exists(config_dir + MODEL_CONFIG):
    print("Downloading config file...")
    wget.download(f"https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/text2sparql/conf/{MODEL_CONFIG}", config_dir)
else:
    print("config file is already exists")

In [ ]:

# this line will print the entire default config of the model
config_path = f"{WORK_DIR}/configs/{MODEL_CONFIG}"
print(config_path)
config = OmegaConf.load(config_path)
print(OmegaConf.to_yaml(config))

Setting Up Data Within Config

The default configuration is missing certain essential fields. We must apply these changes to the configuration:

• config.model.train_ds.filepath: filepath to the train dataset
• config.model.validation_ds.filepath: filepath to the validation dataset
• config.model.test_ds.filepath: filepath to the test dataset
• config.exp_manager.exp_dir: path to the experiment directory

In [ ]:

config.trainer.gpus = GPUS
config.trainer.max_epochs = MAX_EPOCHS

config.model.nemo_path = NEMO_PATH
config.model.batch_size = BATCH_SIZE
config.model.max_seq_length = MAX_SEQ_LENGTH

config.model.language_model.pretrained_model_name = PRETRAINED_BART_MODEL
config.model.encoder_tokenizer.tokenizer_name = TOKENIZER_NAME
config.model.decoder_tokenizer.tokenizer_name = TOKENIZER_NAME

config.model.train_ds.filepath = TRAIN_FILE
config.model.validation_ds.filepath = EVAL_FILE
config.model.test_ds.filepath = TEST_FILE

config.model.optim.lr = LEARNING_RATE

config.exp_manager.exp_dir = WORK_DIR

Building the PyTorch Lightning Trainer

NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem!

Let's first instantiate a Trainer object!

In [ ]:

trainer = pl.Trainer(**config.trainer)

Setting up a NeMo Experiment

NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it!

In [ ]:

exp_dir = exp_manager(trainer, config.get("exp_manager", None))

In [ ]:

nmt_model = NeuralMachineTranslationModel(config.model, trainer=trainer)

Monitoring Training Progress

Optionally, you can create a Tensorboard visualization to monitor training progress.

In [ ]:

# load the TensorBoard notebook extension
try:
    from google import colab
    COLAB_ENV = True
except (ImportError, ModuleNotFoundError):
    COLAB_ENV = False

# Load the TensorBoard notebook extension
if COLAB_ENV:
    %load_ext tensorboard
    %tensorboard --logdir "{exp_dir}"
else:
    print("To use tensorboard, please use this notebook in a Google Colab environment.")

In [ ]:

# start model training
trainer.fit(nmt_model)

Saving and Reloading

We can now save our model as a .nemo package.

In [ ]:

nmt_model.save_to(config.model.nemo_path)

And easily reload it along with its configuration.

In [ ]:

nmt_model = NeuralMachineTranslationModel.restore_from(restore_path=config.model.nemo_path)

Restoring a model does not load our train / validation / test data. Let's load our test set so we can evaluate the model.

In [ ]:

nmt_model.setup_test_data(config.model.validation_ds)

Inference

To see how the model performs, let's run inference on the test data.

In [ ]:

trainer = pl.Trainer(gpus=config.trainer.gpus)
results = trainer.test(nmt_model)

In [ ]:

# save results
with open(config.model.test_ds.filepath, "r") as f:
    lines = f.readlines()

lines[0] = lines[0].strip() + f"\tpredictions\n"
for i, res in enumerate(results[0]["texts"]):
    lines[i + 1] = lines[i + 1].strip() + f"\t{res}\n"

savepath = os.path.join(config.exp_manager.exp_dir, os.path.basename(config.model.test_ds.filepath))
with open(savepath, "w") as f:
    f.writelines(lines)
    print(f"Predictions saved to {savepath}")

Predictions on test easy should look like such:

Sentence: What is the type of 1,1,1-trifluoro-2-chloro-2-bromoethane?
Label: SELECT ?end WHERE { [ 1,1,1-trifluoro-2-chloro-2-bromoethane ] wdt:P31 ?end . }
Predictions: SELECT ?end WHERE { [1,1,1-trifluoro-2-chloro-2-bromoethane] wdt:P31 ?end . }

Aside from spacing, our finetuned network produces nearly identical results!

Entity Resolution

This section contains a practical demonstration of predictions. In order to process a WikiData query, we must also resolve entity names to q-values.

Feel free to skip this section as it does not pertain to NeMo and requires additional files and dependencies.

In [ ]:

! pip install rapidfuzz

# define helper methods for entity resolution
# matching of entities to q-values is easily performed using levenshtein distance
import os
import re
import json
from typing import Dict, Tuple
from urllib.request import Request, urlopen

from rapidfuzz import fuzz, process
from glob import glob

from nemo.collections.nlp.data.data_utils.data_preprocessing import if_exist

base_url = "https://m.meetkai.com/public_datasets/knowledge/"
prefixes = {
    "television_series-5k-preprocessed.json",
    "person-5k-preprocessed.json",
    "movie-5k-preprocessed.json",
    "literary_work-5k-preprocessed.json",
    "chemical-5k-preprocessed.json",
}


def download_entities(infold: str):
    """Downloads text2sparql entity files

    Args:
        infold: save directory path
    """
    os.makedirs(infold, exist_ok=True)

    for prefix in prefixes:
        url = base_url + prefix

        print(f"Downloading: {url}")
        if if_exist(infold, [prefix]):
            print("** Download file already exists, skipping download")
        else:
            req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with open(os.path.join(infold, prefix), "wb") as handle:
                handle.write(urlopen(req, timeout=20).read())


def url_to_qvalue(url: str) -> str:
    """Get q value from Wikidata url
    http://www.wikidata.org/entity/Q494 -> Q494
    """
    return url.split("/")[-1]


def load_entities(data_dir: str) -> Dict:
    assert os.path.isdir(data_dir), f"{data_dir} is not a valid directory."
    fps = glob(os.path.join(data_dir, "*-5k-preprocessed.json"))
    data = []
    for fp in fps:
        with open(fp, "r") as f:
            file_data = json.load(f)
            data.extend(file_data)
    assert data, f"No data was found, please check {data_dir}."

    # expand data
    expanded_data = dict()
    for item in data:
        qvalue = url_to_qvalue(item["thing"])
        for label in item["labels"]:
            expanded_data[label] = qvalue

    return expanded_data


class EntityResolver:
    def __init__(self, data_dir: str, score_cutoff: float = 80.0):
        self.entity_dict = load_entities(data_dir)
        self.choices = list(self.entity_dict.keys())
        self.score_cutoff = score_cutoff

    def resolve_entity(self, entity: str) -> Tuple[str, str, float]:
        """Finds the fuzzy entity match above the cutoff score and returns the Q-value for it.

        Returns:
            (q-value, top_entity_match, simple_levenshtein_score)
        """
        top = process.extractOne(entity, self.choices, scorer=fuzz.ratio, score_cutoff=self.score_cutoff)
        if not top:
            raise ValueError(f"For entity [{entity}], no valid match above cutoff found.")
        return self.lookup(top[0]), top[0], top[1]

    def lookup(self, entity: str):
        return self.entity_dict[entity]

    def resolve(self, text: str) -> str:
        """Replaces all entity matches within a predicted query
        assuming entities are always formatted between two brackets.
        """
        matches = re.findall(r"\[(.*?)\]", text)
        for entity in matches:
            try:
                q_value = self.resolve_entity(entity)[0]
                text = text.replace(f"[{entity}]", f"wd:{q_value}")
            except ValueError as e:
                print(f"WARNING: {e}")
        return text

In [ ]:

# Download entity data
download_entities(DATA_DIR)

# Load the data to the entity resolver
resolver = EntityResolver(data_dir=DATA_DIR)

In [ ]:

# Let's use a cherry picked example prediction as not all queries will (and should) return results
example_prediction = "SELECT ?end WHERE { [1,1,1-trifluoro-2-chloro-2-bromoethane] wdt:P31 ?end . }"
example_query = resolver.resolve(example_prediction)
example_query

In [ ]:

# And query WikiData
import requests
url = "https://query.wikidata.org/sparql"
response = requests.get(url, params = {"format": "json", "query": example_query})
data = response.json()
data

We get 5 results:

• http://www.wikidata.org/entity/Q11173 (chemical compound)
• http://www.wikidata.org/entity/Q12140 (medication)
• http://www.wikidata.org/entity/Q35456 (essential medicine)
• http://www.wikidata.org/entity/Q909194 (inhalational anaesthetic)
• http://www.wikidata.org/entity/Q72941151 (developmental toxicant)

Telling us that 1,1,1-trifluoro-2-chloro-2-bromoethane is the name of a chemical compound used in some medications!

Note: Keep in mind that not all queries will return a result as the data or properties may not exist within WikiData.