single_country_retrain.ipynb

Kernel: Python 3 (ipykernel)

import os
import compress_pickle
import pickle
from deepparse import download_from_public_repository
from deepparse.dataset_container import PickleDatasetContainer
from deepparse.parser import AddressParser
import shutil
from poutyne import set_seeds
import poutyne
import timeit

seed = 42
set_seeds(seed)

# Retrain an Address Parser for Single Country Uses

In this example, we will retrain a pre-trained model to maximize its performance for specific countries (e.g. the UK or Canada).

Retrain a Model

First, to retrain our supervised model, we need parsed address example, as shown in the following figure. Fortunately, we have access to a public dataset of such parsed examples, the Structured Multinational Address Dataset.

For our example, we will focus on UK addresses since we want to parse addresses only from the UK. So let's first download the dataset directly from the public repository using Deepparse download_from_public_repository function.

os.makedirs("dataset")
download_from_public_repository("dataset/data", "", file_extension="zip")

The dataset archive is a zip directory of subdirectories in which each country's data is compressed into an LZMA file (a more aggressive compression algorithm). The dataset public repository offers a script to decompress the LZMA compress dataset zip archive. We will use the basic idea of it to decompress the dataset in the next code cell (the script handles CLI parameters).

# First, let's decompress the archive
archive_root_path = os.path.join("dataset")
archive_path = os.path.join(archive_root_path, "data.zip")

# Unzip the archive
shutil.unpack_archive(archive_path, archive_root_path)

# Delete the archive
os.remove(archive_path)

# The script functions with minor modification to handle argument
# instead or CLI parsed argument


# Function to handle the files paths
def absolute_file_paths(directory):
    """
    Function to get all the absolute paths of files into a directory.
    """
    for dir_path, _, filenames in os.walk(directory):
        for f in filenames:
            if f.endswith(".lzma"):
                yield os.path.abspath(os.path.join(dir_path, f))


# Function to LZMA decompress the files_directory into the path_to_save directory
def lzma_decompress(files_directory, root_path_to_save) -> None:
    """
    Script to decompress the dataset from LZMA compress files into pickled one.
    """
    paths = absolute_file_paths(files_directory)

    for path in paths:
        pickled_data = compress_pickle.load(path, compression="lzma")
        filename = path.split(os.path.sep)[-1].replace(".lzma", ".p")
        file_path = os.path.join(*path.split(os.path.sep)[-4:-1])
        path_to_save = os.path.join(root_path_to_save, file_path)
        os.makedirs(path_to_save, exist_ok=True)
        with open(os.path.join(path_to_save, filename), "wb") as file:
            pickle.dump(pickled_data, file)
        os.remove(path)

# Let's decompress the dataset. It takes several minutes to decompress.

root_dir = os.path.join("dataset", "data")
clean_root_dir = os.path.join(root_dir, "clean_data")
clean_train_directory = os.path.join(clean_root_dir, "train")
clean_test_directory = os.path.join(clean_root_dir, "test")

# We decompress all the dataset
lzma_decompress(root_dir, "dataset")

Now, let's import our train and test datasets into memory to retrain our parser model.

clean_root_dir = os.path.join(root_dir, "clean_data")
clean_train_directory = os.path.join(clean_root_dir, "train")
clean_test_directory = os.path.join(clean_root_dir, "test")

uk_training_data_path = os.path.join(clean_train_directory, "gb.p")
uk_test_data_path = os.path.join(clean_test_directory, "gb.p")

training_container = PickleDatasetContainer(uk_training_data_path)
test_container = PickleDatasetContainer(uk_test_data_path)

We will use the FastText one for our base pre-trained model since it is faster to retrain.

address_parser = AddressParser(model_type="fasttext", device=0)

Loading the embeddings model

But first, let's see what the performance is before retraining.

address_parser.test(test_container, batch_size=256)

_ = address_parser.retrain(
    training_container,
    train_ratio=0.8,
    epochs=1,
    batch_size=32,
    num_workers=2,
    learning_rate=0.001,
    logging_path="./uk_retrain",
    name_of_the_retrain_parser="UKParser",
)

Epoch: 1/5 Train steps: 2500 Val steps: 625 57.82s loss: 0.096870 accuracy: 99.663765 val_loss: 0.105059 val_accuracy: 99.660023
Epoch 1: val_loss improved from inf to 0.10506, saving file to ./uk_faster_retrain/checkpoint_epoch_1.ckpt
Epoch: 2/5 Train steps: 2500 Val steps: 625 58.84s loss: 0.092458 accuracy: 99.677379 val_loss: 0.103238 val_accuracy: 99.672032
Epoch 2: val_loss improved from 0.10506 to 0.10324, saving file to ./uk_faster_retrain/checkpoint_epoch_2.ckpt
Epoch: 3/5 Train steps: 2500 Val steps: 625 58.43s loss: 0.090964 accuracy: 99.683519 val_loss: 0.103035 val_accuracy: 99.673781
Epoch 3: val_loss improved from 0.10324 to 0.10304, saving file to ./uk_faster_retrain/checkpoint_epoch_3.ckpt
Epoch: 4/5 Train steps: 2500 Val steps: 625 58.37s loss: 0.089921 accuracy: 99.685827 val_loss: 0.103027 val_accuracy: 99.673781
Epoch 4: val_loss improved from 0.10304 to 0.10303, saving file to ./uk_faster_retrain/checkpoint_epoch_4.ckpt
Epoch: 5/5 Train steps: 2500 Val steps: 625 58.31s loss: 0.090967 accuracy: 99.684051 val_loss: 0.103027 val_accuracy: 99.673781
Epoch 5: val_loss improved from 0.10303 to 0.10303, saving file to ./uk_faster_retrain/checkpoint_epoch_5.ckpt
Restoring data from ./uk_faster_retrain/checkpoint_epoch_5.ckpt

[{'epoch': 1,
  'time': 57.81894668377936,
  'loss': 0.09687001281384341,
  'accuracy': 99.66376511882362,
  'val_loss': 0.1050588409877063,
  'val_accuracy': 99.66002270373627},
 {'epoch': 2,
  'time': 58.83542291820049,
  'loss': 0.09245755615734887,
  'accuracy': 99.67737932159272,
  'val_loss': 0.10323802052373769,
  'val_accuracy': 99.67203222624765},
 {'epoch': 3,
  'time': 58.427432637661695,
  'loss': 0.090964393501016,
  'accuracy': 99.68351861406266,
  'val_loss': 0.1030354176204891,
  'val_accuracy': 99.67378086061039},
 {'epoch': 4,
  'time': 58.374891674146056,
  'loss': 0.08992147407911884,
  'accuracy': 99.68582667894091,
  'val_loss': 0.10302727049831104,
  'val_accuracy': 99.67378086061039},
 {'epoch': 5,
  'time': 58.309186859056354,
  'loss': 0.09096737359163201,
  'accuracy': 99.68405119942516,
  'val_loss': 0.10302683650516546,
  'val_accuracy': 99.67378086061039}]

address_parser.test(test_container, batch_size=256)

Running test
Test steps: 57 1.74s test_loss: 0.120875 test_accuracy: 99.575062                                                

{'time': 1.7367468271404505,
 'test_loss': 0.12087451704787924,
 'test_accuracy': 99.5750624169449}

To further improve performance, we could train for longer, increase the training dataset size (the actual size of 100,000 addresses), or rework the Seq2Seq hidden sizes. See the retrain interface documentation for all the training parameters.