Toolkit for training/converting LibreTranslate compatible language models 🚂
import random
import os
import hashlib
from net import download

nllb_langs = {
    "af":"afr_Latn",
    "ak":"aka_Latn",
    "am":"amh_Ethi",
    "ar":"arb_Arab",
    "as":"asm_Beng",
    "ay":"ayr_Latn",
    "az":"azj_Latn",
    "bm":"bam_Latn",
    "be":"bel_Cyrl",
    "bn":"ben_Beng",
    "bho":"bho_Deva",
    "bs":"bos_Latn",
    "bg":"bul_Cyrl",
    "ca":"cat_Latn",
    "ceb":"ceb_Latn",
    "cs":"ces_Latn",
    "ckb":"ckb_Arab",
    "tt":"crh_Latn",
    "cy":"cym_Latn",
    "da":"dan_Latn",
    "de":"deu_Latn",
    "el":"ell_Grek",
    "en":"eng_Latn",
    "eo":"epo_Latn",
    "et":"est_Latn",
    "eu":"eus_Latn",
    "ee":"ewe_Latn",
    "fa":"pes_Arab",
    "fi":"fin_Latn",
    "fr":"fra_Latn",
    "gd":"gla_Latn",
    "ga":"gle_Latn",
    "gl":"glg_Latn",
    "gn":"grn_Latn",
    "gu":"guj_Gujr",
    "ht":"hat_Latn",
    "ha":"hau_Latn",
    "he":"heb_Hebr",
    "hi":"hin_Deva",
    "hr":"hrv_Latn",
    "hu":"hun_Latn",
    "hy":"hye_Armn",
    "nl":"nld_Latn",
    "ig":"ibo_Latn",
    "ilo":"ilo_Latn",
    "id":"ind_Latn",
    "is":"isl_Latn",
    "it":"ita_Latn",
    "jv":"jav_Latn",
    "ja":"jpn_Jpan",
    "kn":"kan_Knda",
    "ka":"kat_Geor",
    "kk":"kaz_Cyrl",
    "km":"khm_Khmr",
    "rw":"kin_Latn",
    "ko":"kor_Hang",
    "ku":"kmr_Latn",
    "lo":"lao_Laoo",
    "lv":"lvs_Latn",
    "ln":"lin_Latn",
    "lt":"lit_Latn",
    "lb":"ltz_Latn",
    "lg":"lug_Latn",
    "lus":"lus_Latn",
    "mai":"mai_Deva",
    "ml":"mal_Mlym",
    "mr":"mar_Deva",
    "mk":"mkd_Cyrl",
    "mg":"plt_Latn",
    "mt":"mlt_Latn",
    "mni-Mtei":"mni_Beng",
    "mni":"mni_Beng",
    "mn":"khk_Cyrl",
    "mi":"mri_Latn",
    "ms":"zsm_Latn",
    "my":"mya_Mymr",
    "no":"nno_Latn",
    "ne":"npi_Deva",
    "ny":"nya_Latn",
    "om":"gaz_Latn",
    "or":"ory_Orya",
    "pl":"pol_Latn",
    "pt":"por_Latn",
    "ps":"pbt_Arab",
    "qu":"quy_Latn",
    "ro":"ron_Latn",
    "ru":"rus_Cyrl",
    "sa":"san_Deva",
    "si":"sin_Sinh",
    "sk":"slk_Latn",
    "sl":"slv_Latn",
    "sm":"smo_Latn",
    "sn":"sna_Latn",
    "sd":"snd_Arab",
    "so":"som_Latn",
    "es":"spa_Latn",
    "sq":"als_Latn",
    "sr":"srp_Cyrl",
    "su":"sun_Latn",
    "sv":"swe_Latn",
    "sw":"swh_Latn",
    "ta":"tam_Taml",
    "te":"tel_Telu",
    "tg":"tgk_Cyrl",
    "tl":"tgl_Latn",
    "th":"tha_Thai",
    "ti":"tir_Ethi",
    "ts":"tso_Latn",
    "tk":"tuk_Latn",
    "tr":"tur_Latn",
    "ug":"uig_Arab",
    "uk":"ukr_Cyrl",
    "ur":"urd_Arab",
    "uz":"uzn_Latn",
    "vi":"vie_Latn",
    "xh":"xho_Latn",
    "yi":"ydd_Hebr",
    "yo":"yor_Latn",
    "zh-CN":"zho_Hans",
    "zh":"zho_Hans",
    "zh-TW":"zho_Hant",
    "zu":"zul_Latn",
    "pa":"pan_Guru"
}

def sources_changed(sources, out_dir):
    merge_hash_file = os.path.join(out_dir, "merge-hash.txt")
    sources_hash = hashlib.md5("|".join(sorted([k for k in sources])).encode('utf-8')).hexdigest()

    if os.path.isfile(merge_hash_file):
        with open(merge_hash_file, "r", encoding="utf-8") as f:
            merge_hash = f.readline().strip()
            if merge_hash == sources_hash:
                print("No changes in sources")
                return False

    with open(merge_hash_file, "w", encoding="utf-8") as f:
        f.write(sources_hash)

def get_flores_dataset_path(dataset="dev"):
    if dataset != "dev" and dataset != "devtest":
        print(f"Invalid dataset {dataset} (must be either dev or devtest)")
        exit(1)
    current_dir = os.path.dirname(__file__)
    cache_dir = os.path.join(current_dir, "cache")

    flores_dataset = os.path.join(cache_dir, "flores200_dataset", dataset)
    if not os.path.isdir(flores_dataset):
        os.makedirs(cache_dir, exist_ok=True)
    
        # Download first
        print("Downloading flores200 dataset...")
        fname = os.path.join(cache_dir, "flores200.tar.gz")
        flores_url = "https://tinyurl.com/flores200dataset"
        download(flores_url, cache_dir, basename=os.path.basename(fname))

        import tarfile
        with tarfile.open(fname) as f:
            f.extractall(cache_dir)
        
        if os.path.isfile(fname):
            os.unlink(fname)

        if not os.path.isdir(flores_dataset):
            print(f"Cannot download flores200. Please manually download it from {flores_url} and place it in {cache_dir}")
            exit(1)

    return flores_dataset

def get_flores(lang_code, dataset="dev"):
    flores_dataset = get_flores_dataset_path(dataset)
    source = os.path.join(flores_dataset, nllb_langs[lang_code] + f".{dataset}")

    vs = [line.rstrip('\n') for line in open(source, encoding="utf-8")]
    return vs

def extract_flores_val(src_code, tgt_code, out_dir, dataset="dev"):
    src_f = os.path.join(out_dir, "src-val.txt")
    tgt_f = os.path.join(out_dir, "tgt-val.txt")
    if not os.path.isfile(src_f) or not os.path.isfile(tgt_f):
        src_val = get_flores(src_code, dataset)
        tgt_val = get_flores(tgt_code, dataset)
        with open(src_f, 'w', encoding='utf-8') as f:
            f.write("\n".join(src_val) + "\n")
        print(f"Wrote {src_f}")
        with open(tgt_f, 'w', encoding='utf-8') as f:
            f.write("\n".join(tgt_val) + "\n")
        print(f"Wrote {tgt_f}")
    
def merge_shuffle(sources, out_dir, max_eval_sentences=5000):
    if not sources_changed(sources, out_dir):
        return False

    data = []
    for k in sources:
        source = sources[k]['source']
        target = sources[k]['target']
        
        print(f"Reading {source} - {target}")
        with open(source, "r", encoding="utf-8") as fs, \
             open(target, "r", encoding="utf-8") as ft:
             while True:
                line_s = fs.readline()
                line_t = ft.readline()
                if not line_s or not line_t:
                    break
                data.append((line_s, line_t))
        print(f"New sentence count: {len(data)}")
    
    print("Shuffling")
    random.shuffle(data)

    if len(data) * 0.2 < max_eval_sentences:
        max_eval_sentences = len(data) * 0.2

    print(f"Training size: {len(data) - max_eval_sentences}")
    print(f"Validation size: {max_eval_sentences}")

    # TODO: can this be faster?
    print("Writing sets")
    os.makedirs(out_dir, exist_ok=True)
    count = 0
    with open(os.path.join(out_dir, "src-val.txt"), "w", encoding="utf-8") as fsv, \
        open(os.path.join(out_dir, "tgt-val.txt"), "w", encoding="utf-8") as ftv, \
        open(os.path.join(out_dir, "src-train.txt"), "w", encoding="utf-8") as fst, \
        open(os.path.join(out_dir, "tgt-train.txt"), "w", encoding="utf-8") as ftt:
        for source, target in data:
            if count < max_eval_sentences:
                fsv.write(source)
                ftv.write(target)
            else:
                fst.write(source)
                ftt.write(target)

            count += 1
    
    return True
    

        

data.py
LibreTranslate/Locomotive
import os
import json
import argparse
import ctranslate2
import sentencepiece
from sacrebleu import corpus_bleu
from data import get_flores

parser = argparse.ArgumentParser(description='Evaluate LibreTranslate compatible models')
parser.add_argument('--config',
    type=str,
    default="model-config.json",
    help='Path to model-config.json. Default: %(default)s')
parser.add_argument('--reverse',
    action='store_true',
    help='Reverse the source and target languages in the configuration and data sources. Default: %(default)s')
parser.add_argument('--bleu',
    action="store_true",
    help='Evaluate BLEU score. Default: %(default)s')
parser.add_argument('--flores-id',
    type=int,
    default=None,
    help='Evaluate this flores sentence ID. Default: %(default)s')
parser.add_argument('--tokens',
    action="store_true",
    help='Display tokens rather than words. Default: %(default)s')


args = parser.parse_args()
try:
    with open(args.config) as f:
        config = json.loads(f.read())
    if args.reverse:
        config["from"], config["to"] = config["to"], config["from"]
except Exception as e:
    print(f"Cannot open config file: {e}")
    exit(1)

current_dir = os.path.dirname(__file__)
cache_dir = os.path.join(current_dir, "cache")
model_dirname = f"{config['from']['code']}_{config['to']['code']}-{config['version']}"
run_dir = os.path.join(current_dir, "run", model_dirname)
ct2_model_dir = os.path.join(run_dir, "model")
sp_model = os.path.join(run_dir, "sentencepiece.model")

if not os.path.isdir(ct2_model_dir) or not os.path.isfile(sp_model):
    print(f"The model in {run_dir} is not valid. Did you run train.py first?")
    exit(1)

def translator():
    device = "cuda" if ctranslate2.get_cuda_device_count() > 0 else "cpu"
    model = ctranslate2.Translator(ct2_model_dir, device=device, compute_type="auto")
    tokenizer = sentencepiece.SentencePieceProcessor(sp_model)
    return {"model": model, "tokenizer": tokenizer}

def encode(text, tokenizer):
    return tokenizer.Encode(text, out_type=str)

def decode(tokens):
    if args.tokens:
        return " ".join(tokens)
    else:
        detokenized = "".join(tokens).replace("▁", " ")
        if len(detokenized) > 0 and detokenized[0] == " ":
            detokenized = detokenized[1:]
        return detokenized

data = translator()

if args.bleu or args.flores_id is not None:
    src_text = get_flores(config["from"]["code"], "dev")
    tgt_text = get_flores(config["to"]["code"], "dev")
    
    if args.flores_id is not None:
        src_text = [src_text[args.flores_id]]
        tgt_text = [tgt_text[args.flores_id]]

    translation_obj = data["model"].translate_batch(
        encode(src_text, data["tokenizer"]),
        beam_size=4, # same as argos
        return_scores=False, # speed up
    )

    translated_text = [
        decode(tokens.hypotheses[0])
        for tokens in translation_obj
    ]
    
    bleu_score = round(corpus_bleu(
        translated_text, [[x] for x in tgt_text], tokenize="flores200"
    ).score, 5)

    if args.flores_id is not None:
        print(f"({config['from']['code']})> {src_text[0]}\n(gt)> {tgt_text[0]}\n({config['to']['code']})> {' '.join(translated_text)}")
    else:
        print(f"BLEU score: {bleu_score}")
else:
    # Interactive mode
    print("Starting interactive mode")

    while True:
        try:
            text = input(f"({config['from']['code']})> ")
        except KeyboardInterrupt:
            print("")
            exit(0)

        src_text = [text.rstrip('\n')]
        translation_obj = data["model"].translate_batch(
            encode(src_text, data["tokenizer"]),
            beam_size=4, # same as argos
            return_scores=False, # speed up
        )
        translated_text = [
            decode(tokens.hypotheses[0])
            for tokens in translation_obj
        ]
        print(f"({config['to']['code']})> {translated_text[0]}")
eval.py
LibreTranslate/Locomotive
import argparse
import io
import json
import os
import hashlib

parser = argparse.ArgumentParser(description='Find text in data sources')
parser.add_argument('--config',
    type=str,
    default="model-config.json",
    help='Path to model-config.json. Default: %(default)s')
parser.add_argument('--text',
    type=str,
    help='Text to search. Default: %(default)s')

args = parser.parse_args() 
try:
    with open(args.config) as f:
        config = json.loads(f.read())
except Exception as e:
    print(f"Cannot open config file: {e}")
    exit(1)

current_dir = os.path.dirname(__file__)
cache_dir = os.path.join(current_dir, "cache")
sources = {}

text = args.text.lower()

for s in config['sources']:
    md5 = hashlib.md5(s.encode('utf-8')).hexdigest()
    source_dir = os.path.join(cache_dir, md5)
    if os.path.isdir(source_dir):
        source, target = None, None
        for f in [f.path for f in os.scandir(source_dir) if f.is_file()]:
            if "target" in f.lower():
                target = f
            if "source" in f.lower():
                source = f

        if source is not None and target is not None:
            def scan(file):
                with open(file, 'r', encoding='utf-8') as f:
                    i = 1
                    for line in f:
                        if text in line.lower():
                            print(f"{os.path.basename(s)} ({file}):{i} => {line}")
                        i += 1
            scan(source)
            scan(target)                
        else:
            print(f"Cannot find a source.txt and a target.txt in {s} ({dir}). Skipping...")


find.py
LibreTranslate/Locomotive
import requests
import math
import os
import time
try:
    import queue
except ImportError:
    import Queue as queue
import threading

from urllib3.exceptions import ReadTimeoutError

class AtomicCounter:
    """Atomic, thread-safe counter."""

    def __init__(self, val=0):
        """Initialize a new atomic counter to given initial value (default 0)."""
        self.value = val
        self.lock = threading.Lock()

    def increment(self, val=1):
        """Atomically increment the counter return the new value."""
        with self.lock:
            self.value += val
            return self.value


class NetError(Exception):
    pass

class RangeNotAvailableError(NetError):
    pass

def download(url, destination, progress_callback=None, basename=None, parallel_downloads=16, parallel_chunks_size=10, timeout=30):
    """Download files in parallel (download accelerator)

    Args:
        url (str): URL to download
        destination (str): directory where to download file. If the directory does not exist, it will be created.
        progress_callback (function): an optional callback with one parameter, the download progress percentage.
        parallel_downloads (int): maximum number of parallel downloads if the node supports http range.
        parallel_chunks_size (int): size in MB of chunks for parallel downloads
        timeout (int): seconds before timing out
    Returns:
        str: path to file
    """
    if not os.path.exists(destination):
        os.makedirs(destination, exist_ok=True)

    try:

        download_stream = requests.get(url, timeout=timeout, stream=True)
        headers = download_stream.headers
        
        if basename is not None:
            output_path = os.path.join(destination, basename)
        else:
            output_path = os.path.join(destination, os.path.basename(url))

        # Keep track of download progress (if possible)
        content_length = download_stream.headers.get('content-length')
        total_length = int(content_length) if content_length is not None else None
        downloaded = 0
        chunk_size = int(parallel_chunks_size * 1024 * 1024)
        use_fallback = False
        accept_ranges = headers.get('accept-ranges')

        # Can we do parallel downloads?
        if accept_ranges is not None and accept_ranges.lower() == 'bytes' and total_length is not None and total_length > chunk_size and parallel_downloads > 1:
            num_chunks = int(math.ceil(total_length / float(chunk_size)))
            num_workers = parallel_downloads

            class nonloc:
                completed_chunks = AtomicCounter(0)
                merge_chunks = [False] * num_chunks
                error = None

            def merge():
                current_chunk = 0

                with open(output_path, "wb") as out_file:
                    while current_chunk < num_chunks and nonloc.error is None:
                        if nonloc.merge_chunks[current_chunk]:
                            chunk_file = "%s.part%s" % (output_path, current_chunk)
                            with open(chunk_file, "rb") as fd:
                                out_file.write(fd.read())

                            os.unlink(chunk_file)
                            
                            current_chunk += 1
                        else:
                            time.sleep(0.1)

            def worker():
                while True:
                    task = q.get()
                    part_num, bytes_range = task
                    if bytes_range is None or nonloc.error is not None:
                        q.task_done()
                        break

                    try:
                        # Download chunk
                        res = requests.get(url, stream=True, timeout=timeout, headers={'Range': 'bytes=%s-%s' % bytes_range})
                        if res.status_code == 206:
                            with open("%s.part%s" % (output_path, part_num), 'wb') as fd:
                                bytes_written = 0
                                try:
                                    for chunk in res.iter_content(4096):
                                        bytes_written += fd.write(chunk)
                                except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
                                    raise NetError(str(e))
                                
                                if bytes_written != (bytes_range[1] - bytes_range[0] + 1):
                                    # Process again
                                    q.put((part_num, bytes_range))
                                    return

                            with nonloc.completed_chunks.lock:
                                nonloc.completed_chunks.value += 1

                                if progress_callback is not None:
                                    progress_callback(100.0 * nonloc.completed_chunks.value / num_chunks)
                        
                            nonloc.merge_chunks[part_num] = True
                        else:
                            nonloc.error = RangeNotAvailableError()
                    except NetError as e:
                        time.sleep(5)
                        q.put((part_num, bytes_range))
                    except Exception as e:
                        nonloc.error = e
                    finally:
                        q.task_done()

            q = queue.PriorityQueue()
            threads = []
            for i in range(num_workers):
                t = threading.Thread(target=worker)
                t.start()
                threads.append(t)

            merge_thread = threading.Thread(target=merge)
            merge_thread.start()

            range_start = 0

            for i in range(num_chunks):
                range_end = min(range_start + chunk_size - 1, total_length - 1)
                q.put((i, (range_start, range_end)))
                range_start = range_end + 1

            # block until all tasks are done
            while not all(nonloc.merge_chunks) and nonloc.error is None:
                time.sleep(0.1)

            # stop workers
            for i in range(len(threads)):
                q.put((-1, None))
            for t in threads:
                t.join()
            
            merge_thread.join()

            if nonloc.error is not None:
                if isinstance(nonloc.error, RangeNotAvailableError):
                    use_fallback = True
                else:
                    raise nonloc.error
        else:
            use_fallback = True

        if use_fallback:
            # Single connection, boring download
            with open(output_path, 'wb') as fd:
                for chunk in download_stream.iter_content(4096):
                    downloaded += len(chunk)

                    if progress_callback is not None and total_length is not None:
                        progress_callback((100.0 * float(downloaded) / total_length))

                    fd.write(chunk)
                
    except (requests.exceptions.Timeout, requests.exceptions.ConnectionError, ReadTimeoutError) as e:
        raise NetError(e)

    return output_path
net.py
LibreTranslate/Locomotive
import torch
import sys
import math
from onmt.constants import DefaultTokens
from onmt.transforms import register_transform

# From: https://github.com/OpenNMT/OpenNMT-py
# MIT licensed
# Copyright (c) 2017-Present OpenNMT

def average_models(model_files, output, fp32=False):
    vocab = None
    opt = None
    avg_model = None
    avg_generator = None

    for i, model_file in enumerate(model_files):
        m = torch.load(model_file, map_location='cpu')
        model_weights = m['model']
        generator_weights = m['generator']

        if fp32:
            for k, v in model_weights.items():
                model_weights[k] = v.float()
            for k, v in generator_weights.items():
                generator_weights[k] = v.float()

        if i == 0:
            vocab, opt = m['vocab'], m['opt']
            avg_model = model_weights
            avg_generator = generator_weights
        else:
            for (k, v) in avg_model.items():
                avg_model[k].mul_(i).add_(model_weights[k]).div_(i + 1)

            for (k, v) in avg_generator.items():
                avg_generator[k].mul_(i).add_(generator_weights[k]).div_(i + 1)

    final = {"vocab": vocab, "opt": opt, "optim": None,
             "generator": avg_generator, "model": avg_model}
    
    torch.save(final, output)


def sp_vocab_to_onmt_vocab(sp_vocab, onmt_vocab):
    print(f"Converting {sp_vocab}")
    with open(sp_vocab, 'r', encoding="utf-8") as fin:
        with open(onmt_vocab, 'wb') as fout:
            OMIT = (DefaultTokens.UNK, DefaultTokens.BOS, DefaultTokens.EOS)
            for line in fin:
                w, c = line.rstrip("\n").split(None, 1)
                if w in OMIT:
                    continue
                c = math.exp(float(c)) * 1000000
                c = int(c) + 1
                fout.write(f'{w}\t{c}\n'.encode("utf-8"))
    print(f"Wrote {onmt_vocab}")

onmt_tools.py
LibreTranslate/Locomotive
import requests
import datetime
import json
import os

API_BASE = "https://opus.nlpl.eu/opusapi/"

def opus_datasets():
    r = requests.get(f"{API_BASE}?corpora=True")
    if r.status_code != 200:
        raise Exception(f"Cannot fetch dataset list from opus (status: {r.status_code})")
    res = r.json()
    return res["corpora"]

def get_opus_dataset_url(corpora, from_code, to_code, run_dir):
    current_dir = os.path.dirname(__file__)
    opus_cache = os.path.join(run_dir, "opus_cache.json")

    cache = {}
    if os.path.isfile(opus_cache):
        with open(opus_cache, "r", encoding="utf-8") as f:
            cache = json.loads(f.read())
    
    key = f"{corpora}-{from_code}-{to_code}"
    if key in cache:
        return cache[key]

    r = requests.get(f"{API_BASE}?corpus={corpora}&source={from_code}&target={to_code}&preprocessing=moses&version=latest")
    if r.status_code != 200:
        raise Exception(f"Cannot fetch dataset from opus (status: {r.status_code})")
    res = r.json()

    if not "corpora" in res:
        raise Exception(f"Invalid response from opus: {res}")
    if len(res["corpora"]) == 0:
        raise Exception(f"No corpora is available for {corpora} ({from_code}-{to_code})")
    if len(res["corpora"]) > 1:
        print(f"WARN: Multiple corpora found for {corpora} ({from_code}-{to_code}), using first")
    
    url = res["corpora"][0]["url"]
    cache[key] = url

    with open(opus_cache, "w", encoding="utf-8") as f:
        f.write(json.dumps(cache))

    return url


if __name__ == "__main__":
    print(f"# OPUS datasets")
    print("")
    print(f"Updated: {datetime.date.today()}")
    print("")

    ds = opus_datasets()
    for d in ds:
        print(f" * {d}")

opus.py
LibreTranslate/Locomotive
import argparse
import io
import sys
import json
import os
import argparse
import hashlib
import zipfile
import re
import shutil
import glob
import stanza
import subprocess
from net import download
import requests
import iso639

parser = argparse.ArgumentParser(description='Convert OPUS-MT bilingual models to LibreTranslate compatible models')
parser.add_argument('-s', '--source',
    type=str,
    default="en",
    required=True,
    help='Source language code: %(default)s')
parser.add_argument('-t', '--target',
    type=str,
    default="it",
    required=True,
    help='Target language code: %(default)s')
parser.add_argument('--src-name',
    type=str,
    default="",
    help='Override source name: %(default)s')
parser.add_argument('--tgt-name',
    type=str,
    default="",
    help='Override target name: %(default)s')
parser.add_argument('--model-url',
    type=str,
    default="",
    help='URL to OPUS model: %(default)s')
args = parser.parse_args()

def lang_name_from_code(code):
    lang = iso639.find(code)
    if lang is None:
        print(f"Cannot find source language code: {args.source}")
        exit(1)
    
    name = lang['name']
    if ";" in name:
        name = name[name.index(";")+1:].strip()
    return name
    
if args.src_name:
    src_lang_name = args.src_name
else:
    src_lang_name = lang_name_from_code(args.source)

if args.tgt_name:
    tgt_lang_name = args.tgt_name
else:
    tgt_lang_name = lang_name_from_code(args.target)

model_url = args.model_url
if not model_url:
    # Fetch
    readme_url = f"https://raw.githubusercontent.com/Helsinki-NLP/OPUS-MT-train/master/models/{args.source}-{args.target}/README.md"
    print(f"Fetching {readme_url}")
    r = requests.get(readme_url)
    readme = r.content.decode("utf-8")
    matches = None
    found_sentencepiece = False
    for line in readme.split("\n"):
        if "sentencepiece" in line.lower():
            found_sentencepiece = True
            continue

        if found_sentencepiece:
            matches = re.match(".*download: \[[^\]]+\]\((http[^)]+)\)", line)
            if matches:
                break

    if matches is None:
        print("Cannot find opus model URL. Please provide it manually via --model")
        exit(1)
    model_url = matches[1]

print(f"Model URL: {model_url}")
version = "1.0"

print(f"{src_lang_name} --> {tgt_lang_name} ({version})")

current_dir = os.path.dirname(__file__)
cache_dir = os.path.join(current_dir, "cache")
os.makedirs(cache_dir, exist_ok=True)

model_dirname = f"{args.source}_{args.target}-opus_{version}"
run_dir = os.path.join(current_dir, "run", model_dirname)
stanza_dir = os.path.join(run_dir, "stanza")

md5 = hashlib.md5(model_url.encode('utf-8')).hexdigest()
model_path = os.path.join(cache_dir, md5) 
zip_path = model_path + ".zip"

if not os.path.isdir(model_path):
    def print_progress(progress):
        print(f"\r{model_url} [{int(progress)}%]     ", end='\r')

    download(model_url, cache_dir, progress_callback=print_progress, basename=os.path.basename(zip_path))

    os.makedirs(model_path, exist_ok=True)
    print(f"Extracting {zip_path} to {model_path}")
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(model_path)

    os.unlink(zip_path)

# Identify pieces
spm_models = glob.glob(os.path.join(model_path, "*.spm"))
sp_model = None
for spm in spm_models:
    if "source" in spm.lower():
        sp_model = spm
        break

if sp_model is None:
    print("Cannot find SentencePiece source model")
    exit(1)

npz_models = glob.glob(os.path.join(model_path, "*.npz"))
npz_model = None

# Case with a single model
if len(npz_models) == 1:
    npz_model = npz_models[0]

# Search for "best" string
if npz_model is None:
    for nm in npz_models:
        if "best" in nm.lower():
            npz_model = nm
            break

# Search for largest size
if npz_model is None:
    largest = float('-inf')
    for nm in npz_models:
        size = os.path.getsize(nm)
        if size > largest:
            largest = size
            npz_model = nm

if npz_model is None:
    print("Cannot find npz model")
    exit(1)


vocabs = glob.glob(os.path.join(model_path, "*.yml"))
vocab_file = None
# Case with a single model
for v in vocabs:
    if "opus" in v.lower() and not "decoder" in v.lower():
        vocab_file = v
        break

if vocab_file is None:
    print("Cannot find vocab file")
    exit(1)

print(f"SentencePiece model: {sp_model}")
print(f"OPUS model: {npz_model}")
print(f"Vocab: {vocab_file}")

if os.path.isdir(run_dir):
    shutil.rmtree(run_dir)
os.makedirs(run_dir, exist_ok=True)

# Copy sp model
sp_model_path = os.path.join(run_dir, "sentencepiece.model")
shutil.copy(sp_model, sp_model_path)

stanza_lang_code = args.source
if not os.path.isdir(os.path.join(stanza_dir, stanza_lang_code)):
    while True:
        try:
            os.makedirs(stanza_dir, exist_ok=True)
            stanza.download(stanza_lang_code, dir=stanza_dir, processors="tokenize")
            break
        except Exception as e:
            if stanza_lang_code != "en":
                print(f'Could not locate stanza model for "{stanza_lang_code}", we will use "en" instead. Note this might not work well.')
                stanza_lang_code = "en"
            else:
                print(f'Cannot download stanza model: {str(e)}')
                exit(1)

# Quantize
ct2_model_dir = os.path.join(run_dir, "model")
if os.path.isdir(ct2_model_dir):
    shutil.rmtree(ct2_model_dir)

print("Converting to ctranslate2")
subprocess.run([
        "ct2-marian-converter",
        "--model_path",
        npz_model,
        "--vocab_paths",
        vocab_file,
        "--output_dir",
        ct2_model_dir,
        "--quantization",
        "int8"])

# Package
readme = f"""# {src_lang_name} - {tgt_lang_name} version {version}

Authors: Jörg Tiedemann and Santhosh Thottingal
Title: "OPUS-MT — Building open translation services for the World"
Book Title: Proceedings of the 22nd Annual Conference of the European Association for Machine Translation (EAMT)
Year: 2020
Location: Lisbon, Portugal

The original OPUS model from which this packaged model is derived is licensed CC-BY 4.0
"""

metadata = {
    "package_version": version,
    "argos_version": "1.8.2",
    "from_code": args.source,
    "from_name": src_lang_name,
    "to_code": args.target,
    "to_name": tgt_lang_name,
}

readme_file = os.path.join(run_dir, "README.md")

with open(readme_file, "w", encoding="utf-8") as f:
    f.write(readme)

metadata_file = os.path.join(run_dir, "metadata.json")
with open(metadata_file, "w", encoding="utf-8") as f:
    f.write(json.dumps(metadata))


package_file = os.path.join(run_dir, f"translate-{args.source}_{args.target}-{version.replace('.', '_')}.argosmodel")
if os.path.isfile(package_file):
    os.unlink(package_file)

print(f"Writing {package_file}")
with zipfile.ZipFile(package_file, 'w', compression=zipfile.ZIP_STORED) as zipf:
    def add_file(f):
        zipf.write(f, arcname=os.path.basename(f))

    def add_folder(f):
        for root, dirs, files in os.walk(f):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, f)
                zipf.write(file_path, arcname=os.path.join(os.path.basename(f), arcname))

    add_file(readme_file)
    add_file(metadata_file)
    add_file(sp_model_path)
    add_folder(ct2_model_dir)
    add_folder(stanza_dir)

# Write config file
config = {
    "from": {
        "name": src_lang_name,
        "code": args.source
    },
    "to": {
        "name": tgt_lang_name,
        "code": args.target
    },
    "version": f"opus_{version}",
    "sources": [
    ]
}
config_file = os.path.join(run_dir, "config.json")
with open(config_file, "w", encoding="utf-8") as f:
    f.write(json.dumps(config))
print(f"Wrote {config_file}")

print("Done!")
opus_mt_convert.py
LibreTranslate/Locomotive
import argparse
import io
import sys
import json
import os
import argparse
import hashlib
import zipfile
import shutil
import glob
import yaml
import subprocess
import stanza
import re
import zipfile
import ctranslate2
from opus import get_opus_dataset_url
from net import download
from data import sources_changed, get_flores, extract_flores_val
import sentencepiece as spm
from onmt_tools import average_models, sp_vocab_to_onmt_vocab

parser = argparse.ArgumentParser(description='Train LibreTranslate compatible models')
parser.add_argument('--config',
    type=str,
    default="model-config.json",
    help='Path to model-config.json. Default: %(default)s')
parser.add_argument('--reverse',
    action='store_true',
    help='Reverse the source and target languages in the configuration and data sources. Default: %(default)s')
parser.add_argument('--rerun',
    action='store_true',
    help='Rerun the training from scratch. Default: %(default)s')
parser.add_argument('--tensorboard',
    action='store_true',
    help='Run tensorboard during training. Default: %(default)s')
parser.add_argument('--toy',
    action='store_true',
    help='Train a toy model (useful for testing). Default: %(default)s')
parser.add_argument('--inflight',
    action='store_true',
    help='While training is in progress on a separate process, you can launch another instance of train.py with this flag turned on to build a model from the last available checkpoints rather that waiting until the end. Default: %(default)s')

args = parser.parse_args() 
try:
    with open(args.config) as f:
        config = json.loads(f.read())
    
    if args.reverse:
        config["from"], config["to"] = config["to"], config["from"]
except Exception as e:
    print(f"Cannot open config file: {e}")
    exit(1)

print(f"Training {config['from']['name']} --> {config['to']['name']} ({config['version']})")
print(f"Sources: {len(config['sources'])}")

metadata = {
    "package_version": config['version'],
    "argos_version": "1.8.2",
    "from_code": config['from']['code'],
    "from_name": config['from']['name'],
    "to_code": config['to']['code'],
    "to_name": config['to']['name'],
}
readme = f"# {config['from']['name']} - {config['to']['name']} version {config['version']}"

current_dir = os.path.dirname(__file__)
cache_dir = os.path.join(current_dir, "cache")
model_dirname = f"{config['from']['code']}_{config['to']['code']}-{config['version']}"
run_dir = os.path.join(current_dir, "run", model_dirname)
onmt_dir = os.path.join(run_dir, "opennmt")
stanza_dir = os.path.join(run_dir, "stanza")
rel_run_dir = f"run/{model_dirname}"
rel_onmt_dir = f"{rel_run_dir}/opennmt"
os.makedirs(cache_dir, exist_ok=True)

if args.rerun and os.path.isdir(run_dir):
    shutil.rmtree(run_dir)
os.makedirs(run_dir, exist_ok=True)

sources = {}

for s in config['sources']:
    weight = 1
    if isinstance(s, dict):
        if not "source" in s:
            print("Malformed source: {s}. A 'source' key is required.")
        weight = s.get("weight", 1)
        s = s["source"]

    md5 = hashlib.md5(s.encode('utf-8')).hexdigest()
    
    def add_source_from(dir):
        source, target = None, None
        skip_reverse = False
        for f in [f.path for f in os.scandir(dir) if f.is_file()]:
            if "target" in f.lower():
                target = f
            elif f.lower().endswith(f".{config['to']['code']}"):
                target = f
                skip_reverse = True
            
            if "source" in f.lower():
                source = f
            elif f.lower().endswith(f".{config['from']['code']}"):
                source = f
                skip_reverse = True

        if source is not None and target is not None:
            if args.reverse and not skip_reverse:
                source, target = target, source
            sources[s] = {
                'source': source,
                'target': target,
                'hash': md5,
                'weight': weight
            }
        else:
            print(f"Cannot find a source.txt and a target.txt in {s} ({dir}). Exiting...")
            exit(1)

    if s.lower().startswith("file://"):
        add_source_from(s[7:])
    else:
        if s.lower().startswith("opus://"):
            try:
                s = get_opus_dataset_url(s[7:], config["from"]["code"], config["to"]["code"], run_dir)
            except Exception as e:
                print(e)
                exit(1)

        # Network/OPUS URL
        dataset_path = os.path.join(cache_dir, md5)
        zip_path = dataset_path + ".zip"

        # Download first?
        if not os.path.isdir(dataset_path):
            def download_source():
                def print_progress(progress):
                    print(f"\r{s} [{int(progress)}%]     ", end='\r')
                
                download(s, cache_dir, progress_callback=print_progress, basename=os.path.basename(zip_path))
                print()

            if not os.path.isfile(zip_path):
                download_source()
            else:
                # Quick check
                try:
                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                        pass
                except:
                    print(f"Corrupted .zip file, redownloading {zip_path}")
                    os.unlink(zip_path)
                    download_source()

            os.makedirs(dataset_path, exist_ok=True)
            print(f"Extracting {zip_path} to {dataset_path}")
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(dataset_path)
            
            os.unlink(zip_path)
        
            subfolders = [ f.path for f in os.scandir(dataset_path) if f.is_dir()]
            if len(subfolders) == 1:
                # Move files from subfolder
                for f in [f.path for f in os.scandir(subfolders[0]) if f.is_file()]:
                    shutil.move(f, dataset_path)
                
                shutil.rmtree(subfolders[0])
        
        add_source_from(dataset_path)

for k in sources:
    print(f" - {k} ({sources[k]['hash']})")

stanza_lang_code = config['from']['code']
if not os.path.isdir(os.path.join(stanza_dir, stanza_lang_code)):
    while True:
        try:
            os.makedirs(stanza_dir, exist_ok=True)
            stanza.download(stanza_lang_code, dir=stanza_dir, processors="tokenize")
            break
        except Exception as e:
            if stanza_lang_code != "en":
                print(f'Could not locate stanza model for "{stanza_lang_code}", we will use "en" instead. Note this might not work well.')
                stanza_lang_code = "en"
            else:
                print(f'Cannot download stanza model: {str(e)}')
                exit(1)

extract_flores_val(config['from']['code'], config['to']['code'], run_dir, dataset="devtest")
changed = sources_changed(sources, run_dir)

sp_model_path = os.path.join(run_dir, "sentencepiece.model")
if not os.path.isfile(sp_model_path) or changed:
    while True:
        try:
            spm.SentencePieceTrainer.train(input=[sources[s]["source"] for s in sources] + [sources[s]["target"] for s in sources], 
                                            model_prefix=f"{run_dir}/sentencepiece", vocab_size=config.get('vocab_size', 50000),
                                            character_coverage=config.get('character_coverage', 1.0),
                                            input_sentence_size=config.get('input_sentence_size', 1000000),
                                            shuffle_input_sentence=True)
            break
        except Exception as e:
            err = str(e)
            if "Vocabulary size too high" in err:
                matches = re.match(".*Please set it to a value <= (\d+)", err)
                if matches is not None:
                    config["vocab_size"] = int(matches.group(1))
                    print(f"WARNING: vocabulary size too high, reducing to {matches.group(1)}")
                else:
                    print(err)
                    exit(1)
            else:
                print(err)
                exit(1)

os.makedirs(onmt_dir, exist_ok=True)

transforms = ['sentencepiece', 'filtertoolong']
onmt_config = {
    'save_data': rel_onmt_dir,
    'src_vocab': f"{rel_onmt_dir}/openmt.vocab",
    'tgt_vocab': f"{rel_onmt_dir}/openmt.vocab",
    'src_vocab_size': config.get('vocab_size', 50000),
    'tgt_vocab_size': config.get('vocab_size', 50000),
    'share_vocab': True, 
    'data': {
        'valid': {
            'path_src': f'{rel_run_dir}/src-val.txt',
            'path_tgt': f'{rel_run_dir}/tgt-val.txt', 
            'transforms': transforms
        }
    }, 
    'src_subword_type': 'sentencepiece',
    'tgt_subword_type': 'sentencepiece',
    'src_onmttok_kwargs': {
        'mode': 'none',
        'lang': config['from']['code'],
    },
    'tgt_onmttok_kwargs': {
        'mode': 'none',
        'lang': config['to']['code'],
    },
    'src_subword_model': f'{rel_run_dir}/sentencepiece.model', 
    'tgt_subword_model': f'{rel_run_dir}/sentencepiece.model', 
    'src_subword_nbest': 1, 
    'src_subword_alpha': 0.0, 
    'tgt_subword_nbest': 1, 
    'tgt_subword_alpha': 0.0, 
    'src_seq_length': 150, 
    'tgt_seq_length': 150, 
    'skip_empty_level': 'silent', 
    'save_model': f'{rel_onmt_dir}/openmt.model', 
    'save_checkpoint_steps': 1000, 
    'valid_steps': 2500, 
    'train_steps': 50000, 
    'early_stopping': 4, 
    'bucket_size': 262144, 
    'world_size': 1, 
    'gpu_ranks': [0], 
    'batch_type': 'tokens', 
    'queue_size': 10000,
    'batch_size': 8192, 
    'valid_batch_size': 128,
    'max_generator_batches': 2, 
    'accum_count': [4], 
    'accum_steps': [0], 
    'model_dtype': 'fp16', 
    'optim': 'adam', 
    'learning_rate': 2, 
    'warmup_steps': 8000, 
    'decay_method': 'noam', 
    'adam_beta2': 0.998, 
    'max_grad_norm': 0, 
    'label_smoothing': 0.1, 
    'param_init': 0, 
    'param_init_glorot': True, 
    'normalization': 'tokens', 
    'encoder_type': 'transformer', 
    'decoder_type': 'transformer', 
    'position_encoding': False,
    'max_relative_positions': 20,
    'averaging_decay': 0.0001,
    'enc_layers': 6, 
    'dec_layers': 6,
    'heads': 8,
    'hidden_size': 512, 
    'rnn_size': 512,
    'word_vec_size': 512, 
    'transformer_ff': 2048,
    'dropout_steps': [0],
    'dropout': [0.1],
    'attention_dropout': [0.1],
    'share_decoder_embeddings': True,
    'share_embeddings': True
}

# Populate data sources
for s in sources:
    onmt_config['data'][sources[s]['hash']] = {
        'path_src': sources[s]['source'], 
        'path_tgt': sources[s]['target'], 
        'weight': sources[s]['weight'],
        'transforms': transforms
    }

no_gpu = ctranslate2.get_cuda_device_count() == 0
if sys.platform == 'darwin' or no_gpu:
    # CPU
    del onmt_config['gpu_ranks']

if args.toy:
    toy_config = {
        'valid_steps': 100, 
        'train_steps': 200, 
        'save_checkpoint_steps': 100
    }
    for k in toy_config:
        onmt_config[k] = toy_config[k]

# Config defined overrides
for k in onmt_config:
    if k in config:
        onmt_config[k] = config[k]

onmt_config_path = os.path.join(run_dir, "config.yml")
with open(onmt_config_path, "w", encoding="utf-8") as f:
    f.write(yaml.dump(onmt_config))
    print(f"Wrote {onmt_config_path}")

sp_vocab_file = os.path.join(run_dir, "sentencepiece.vocab")
onmt_vocab_file = os.path.join(onmt_dir, "openmt.vocab")
if changed and os.path.isfile(onmt_vocab_file):
    os.unlink(onmt_vocab_file)
    
if not os.path.isfile(onmt_vocab_file):
    #subprocess.run(["onmt_build_vocab", "-config", onmt_config_path, "-n_sample", "-1", "-num_threads", str(os.cpu_count())])
    sp_vocab_to_onmt_vocab(sp_vocab_file, onmt_vocab_file)

last_checkpoint = os.path.join(onmt_dir, os.path.basename(onmt_config["save_model"]) + f'_step_{onmt_config["train_steps"]}.pt')
if (not (os.path.isfile(last_checkpoint) or args.inflight)) or changed:
    cmd = ["onmt_train", "-config", onmt_config_path]
    if args.tensorboard:
        print("Launching tensorboard")

        from tensorboard import program
        import webbrowser
        import mimetypes

        log_dir = os.path.join(onmt_dir, "logs")
        
        # Allow tensorboard to run on Windows due to mimetypes bug: https://github.com/microsoft/vscode-python/pull/16203
        mimetypes.add_type("application/javascript", ".js")

        tb = program.TensorBoard()
        tb.configure(argv=[None, '--logdir', os.path.abspath(log_dir)])
        url = tb.launch()
        print(f"Tensorboard URL: {url}")
        webbrowser.open(url)

        cmd += ["--tensorboard", "--tensorboard_log_dir", log_dir]
    
    # Resume?
    checkpoints = sorted(glob.glob(os.path.join(onmt_dir, "*.pt")))
    if len(checkpoints) > 0 and not changed:
        print(f"Resuming from {checkpoints[-1]}")
        cmd += ["--train_from", checkpoints[-1]]

    subprocess.run(cmd)

# Average
average_checkpoint = os.path.join(run_dir, "averaged.pt")
checkpoints = sorted(glob.glob(os.path.join(onmt_dir, "*.pt")))
print(f"Total checkpoints: {len(checkpoints)}")

if len(checkpoints) == 0:
    print("Something went wrong, looks like onmt_train failed?")
    exit(1)

if os.path.isfile(average_checkpoint):
    os.unlink(average_checkpoint)

if len(checkpoints) == 1 or args.inflight:
    print("Single checkpoint")
    shutil.copy(checkpoints[-1], average_checkpoint)
else:
    if config.get('avg_checkpoints', 1) == 1:
        print("Averaging 1 model")
        shutil.copy(checkpoints[-1], average_checkpoint)
    else:
        avg_num = min(config.get('avg_checkpoints', 1), len(checkpoints))
        print(f"Averaging {avg_num} models")
        average_models(checkpoints[-avg_num:], average_checkpoint)

# Quantize
ct2_model_dir = os.path.join(run_dir, "model")
if os.path.isdir(ct2_model_dir):
    shutil.rmtree(ct2_model_dir)

print("Converting to ctranslate2")
subprocess.run([
        "ct2-opennmt-py-converter",
        "--model_path",
        average_checkpoint,
        "--output_dir",
        ct2_model_dir,
        "--quantization",
        "int8"])

# Package
readme_file = os.path.join(run_dir, "README.md")
with open(readme_file, "w", encoding="utf-8") as f:
    f.write(readme)

metadata_file = os.path.join(run_dir, "metadata.json")
with open(metadata_file, "w", encoding="utf-8") as f:
    f.write(json.dumps(metadata))

package_file = os.path.join(run_dir, f"translate-{config['from']['code']}_{config['to']['code']}-{config['version'].replace('.', '_')}.argosmodel")
if os.path.isfile(package_file):
    os.unlink(package_file)

print(f"Writing {package_file}")
with zipfile.ZipFile(package_file, 'w', compression=zipfile.ZIP_STORED) as zipf:
    def add_file(f):
        zipf.write(f, arcname=os.path.basename(f))

    def add_folder(f):
        for root, dirs, files in os.walk(f):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, f)
                zipf.write(file_path, arcname=os.path.join(os.path.basename(f), arcname))

    add_file(readme_file)
    add_file(metadata_file)
    add_file(sp_model_path)
    add_folder(ct2_model_dir)
    add_folder(stanza_dir)

print("Done!")

train.py
LibreTranslate/Locomotive