chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,240 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import argparse
|
||||
import glob
|
||||
from subprocess import check_call
|
||||
|
||||
try:
|
||||
import faiss
|
||||
|
||||
has_faiss = True
|
||||
except ImportError:
|
||||
has_faiss = False
|
||||
import numpy as np
|
||||
|
||||
|
||||
GB = 1024 * 1024 * 1024
|
||||
|
||||
|
||||
def call(cmd):
|
||||
print(cmd)
|
||||
check_call(cmd, shell=True)
|
||||
|
||||
|
||||
def get_batches(directory, lang, prefix="all_avg_pool"):
|
||||
print(f"Finding in {directory}/{prefix}.{lang}*")
|
||||
files = glob.glob(f"{directory}/{prefix}.{lang}*")
|
||||
emb_files = []
|
||||
txt_files = []
|
||||
for emb_fi in files:
|
||||
emb_files.append(emb_fi)
|
||||
txt_fi = emb_fi.replace(prefix, "sentences")
|
||||
txt_files.append(txt_fi)
|
||||
return emb_files, txt_files
|
||||
|
||||
|
||||
def load_batch(emb_file, dim):
|
||||
embeddings = np.fromfile(emb_file, dtype=np.float32)
|
||||
num_rows = int(embeddings.shape[0] / dim)
|
||||
embeddings = embeddings.reshape((num_rows, dim))
|
||||
faiss.normalize_L2(embeddings)
|
||||
return embeddings
|
||||
|
||||
|
||||
def knnGPU_sharded(x_batches_f, y_batches_f, dim, k, direction="x2y"):
|
||||
if not has_faiss:
|
||||
raise ImportError("Please install Faiss")
|
||||
sims = []
|
||||
inds = []
|
||||
xfrom = 0
|
||||
xto = 0
|
||||
for x_batch_f in x_batches_f:
|
||||
yfrom = 0
|
||||
yto = 0
|
||||
x_batch = load_batch(x_batch_f, dim)
|
||||
xto = xfrom + x_batch.shape[0]
|
||||
bsims, binds = [], []
|
||||
for y_batch_f in y_batches_f:
|
||||
y_batch = load_batch(y_batch_f, dim)
|
||||
neighbor_size = min(k, y_batch.shape[0])
|
||||
yto = yfrom + y_batch.shape[0]
|
||||
print("{}-{} -> {}-{}".format(xfrom, xto, yfrom, yto))
|
||||
idx = faiss.IndexFlatIP(dim)
|
||||
idx = faiss.index_cpu_to_all_gpus(idx)
|
||||
idx.add(y_batch)
|
||||
bsim, bind = idx.search(x_batch, neighbor_size)
|
||||
|
||||
bsims.append(bsim)
|
||||
binds.append(bind + yfrom)
|
||||
yfrom += y_batch.shape[0]
|
||||
del idx
|
||||
del y_batch
|
||||
bsims = np.concatenate(bsims, axis=1)
|
||||
binds = np.concatenate(binds, axis=1)
|
||||
aux = np.argsort(-bsims, axis=1)
|
||||
sim_batch = np.zeros((x_batch.shape[0], k), dtype=np.float32)
|
||||
ind_batch = np.zeros((x_batch.shape[0], k), dtype=np.int64)
|
||||
for i in range(x_batch.shape[0]):
|
||||
for j in range(k):
|
||||
sim_batch[i, j] = bsims[i, aux[i, j]]
|
||||
ind_batch[i, j] = binds[i, aux[i, j]]
|
||||
sims.append(sim_batch)
|
||||
inds.append(ind_batch)
|
||||
xfrom += x_batch.shape[0]
|
||||
del x_batch
|
||||
sim = np.concatenate(sims, axis=0)
|
||||
ind = np.concatenate(inds, axis=0)
|
||||
return sim, ind
|
||||
|
||||
|
||||
def score(sim, fwd_mean, bwd_mean, margin):
|
||||
return margin(sim, (fwd_mean + bwd_mean) / 2)
|
||||
|
||||
|
||||
def score_candidates(
|
||||
sim_mat, candidate_inds, fwd_mean, bwd_mean, margin, verbose=False
|
||||
):
|
||||
print(" - scoring {:d} candidates".format(sim_mat.shape[0]))
|
||||
scores = np.zeros(candidate_inds.shape)
|
||||
for i in range(scores.shape[0]):
|
||||
for j in range(scores.shape[1]):
|
||||
k = int(candidate_inds[i, j])
|
||||
scores[i, j] = score(sim_mat[i, j], fwd_mean[i], bwd_mean[k], margin)
|
||||
return scores
|
||||
|
||||
|
||||
def load_text(files):
|
||||
all_sentences = []
|
||||
for fi in files:
|
||||
with open(fi) as sentence_fi:
|
||||
for line in sentence_fi:
|
||||
all_sentences.append(line.strip())
|
||||
print(f"Read {len(all_sentences)} sentences")
|
||||
return all_sentences
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Mine bitext")
|
||||
parser.add_argument("--src-lang", help="Source language")
|
||||
parser.add_argument("--tgt-lang", help="Target language")
|
||||
parser.add_argument(
|
||||
"--dict-path", help="Path to dictionary file", default="dict.txt"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spm-path", help="Path to SPM model file", default="sentence.bpe.model"
|
||||
)
|
||||
parser.add_argument("--dim", type=int, default=1024, help="Embedding dimension")
|
||||
parser.add_argument("--mem", type=int, default=5, help="Memory in GB")
|
||||
parser.add_argument("--src-dir", help="Source directory")
|
||||
parser.add_argument("--tgt-dir", help="Target directory")
|
||||
parser.add_argument("--output", help="Output path")
|
||||
parser.add_argument(
|
||||
"--neighborhood", type=int, default=4, help="Embedding dimension"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--threshold", type=float, default=1.06, help="Threshold on mined bitext"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--valid-size",
|
||||
type=int,
|
||||
default=2000,
|
||||
help="Number of sentences used for validation set",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-count",
|
||||
type=int,
|
||||
default=50000,
|
||||
help="Min num sentences used for each language",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
x_batches_f, x_sents_f = get_batches(args.src_dir, args.src_lang)
|
||||
y_batches_f, y_sents_f = get_batches(args.tgt_dir, args.tgt_lang)
|
||||
margin = lambda a, b: a / b
|
||||
y2x_sim, y2x_ind = knnGPU_sharded(
|
||||
y_batches_f, x_batches_f, args.dim, args.neighborhood, direction="y2x"
|
||||
)
|
||||
x2y_sim, x2y_ind = knnGPU_sharded(
|
||||
x_batches_f, y_batches_f, args.dim, args.neighborhood, direction="x2y"
|
||||
)
|
||||
|
||||
x2y_mean = x2y_sim.mean(axis=1)
|
||||
y2x_mean = y2x_sim.mean(axis=1)
|
||||
fwd_scores = score_candidates(x2y_sim, x2y_ind, x2y_mean, y2x_mean, margin)
|
||||
bwd_scores = score_candidates(y2x_sim, y2x_ind, y2x_mean, x2y_mean, margin)
|
||||
fwd_best = x2y_ind[np.arange(x2y_sim.shape[0]), fwd_scores.argmax(axis=1)]
|
||||
bwd_best = y2x_ind[np.arange(y2x_sim.shape[0]), bwd_scores.argmax(axis=1)]
|
||||
indices = np.stack(
|
||||
(
|
||||
np.concatenate((np.arange(x2y_ind.shape[0]), bwd_best)),
|
||||
np.concatenate((fwd_best, np.arange(y2x_ind.shape[0]))),
|
||||
),
|
||||
axis=1,
|
||||
)
|
||||
scores = np.concatenate((fwd_scores.max(axis=1), bwd_scores.max(axis=1)))
|
||||
|
||||
x_sentences = load_text(x_sents_f)
|
||||
y_sentences = load_text(y_sents_f)
|
||||
|
||||
threshold = args.threshold
|
||||
min_count = args.min_count
|
||||
seen_src, seen_trg = set(), set()
|
||||
directory = args.output
|
||||
call(f"mkdir -p {directory}")
|
||||
src_out = open(
|
||||
f"{directory}/all.{args.src_lang}",
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
errors="surrogateescape",
|
||||
)
|
||||
tgt_out = open(
|
||||
f"{directory}/all.{args.tgt_lang}",
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
errors="surrogateescape",
|
||||
)
|
||||
scores_out = open(
|
||||
f"{directory}/all.scores", mode="w", encoding="utf-8", errors="surrogateescape"
|
||||
)
|
||||
count = 0
|
||||
for i in np.argsort(-scores):
|
||||
src_ind, trg_ind = indices[i]
|
||||
if src_ind not in seen_src and trg_ind not in seen_trg:
|
||||
seen_src.add(src_ind)
|
||||
seen_trg.add(trg_ind)
|
||||
if scores[i] > threshold or count < min_count:
|
||||
if x_sentences[src_ind]:
|
||||
print(scores[i], file=scores_out)
|
||||
print(x_sentences[src_ind], file=src_out)
|
||||
print(y_sentences[trg_ind], file=tgt_out)
|
||||
count += 1
|
||||
else:
|
||||
print(f"Ignoring sentence: {x_sentences[src_ind]}")
|
||||
src_out.close()
|
||||
tgt_out.close()
|
||||
scores_out.close()
|
||||
|
||||
print(f"Found {count} pairs for threshold={threshold}")
|
||||
with open(f"{directory}/all.{args.src_lang}") as all_s, open(
|
||||
f"{directory}/all.{args.tgt_lang}"
|
||||
) as all_t, open(f"{directory}/valid.{args.src_lang}", "w") as valid_s, open(
|
||||
f"{directory}/valid.{args.tgt_lang}", "w"
|
||||
) as valid_t, open(
|
||||
f"{directory}/train.{args.src_lang}", "w"
|
||||
) as train_s, open(
|
||||
f"{directory}/train.{args.tgt_lang}", "w"
|
||||
) as train_t:
|
||||
count = 0
|
||||
for s_line, t_line in zip(all_s, all_t):
|
||||
s_line = s_line.split("\t")[1]
|
||||
t_line = t_line.split("\t")[1]
|
||||
if count >= args.valid_size:
|
||||
train_s.write(s_line)
|
||||
train_t.write(t_line)
|
||||
else:
|
||||
valid_s.write(s_line)
|
||||
valid_t.write(t_line)
|
||||
count += 1
|
||||
@@ -0,0 +1,103 @@
|
||||
#!/bin/bash
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
source_lang=kk_KZ
|
||||
target_lang=en_XX
|
||||
MODEL=criss_checkpoints/criss.3rd.pt
|
||||
SPM=criss_checkpoints/sentence.bpe.model
|
||||
SPLIT=test
|
||||
LANG_DICT=criss_checkpoints/lang_dict.txt
|
||||
SPM_ENCODE=flores/scripts/spm_encode.py
|
||||
SAVE_ENCODER=save_encoder.py
|
||||
ENCODER_SAVE_ROOT=sentence_embeddings/$MODEL
|
||||
DICT=criss_checkpoints/dict.txt
|
||||
THRESHOLD=1.02
|
||||
MIN_COUNT=500
|
||||
|
||||
DATA_DIR=data_tmp
|
||||
SAVE_DIR=mining/${source_lang}_${target_lang}_mined
|
||||
ENCODER_SAVE_DIR=${ENCODER_SAVE_ROOT}/${source_lang}-${target_lang}
|
||||
INPUT_DIR=$DATA_DIR/${source_lang}-${target_lang}-tatoeba
|
||||
|
||||
mkdir -p $ENCODER_SAVE_DIR/${target_lang}
|
||||
mkdir -p $ENCODER_SAVE_DIR/${source_lang}
|
||||
mkdir -p $SAVE_DIR
|
||||
|
||||
## Save encoder outputs
|
||||
|
||||
# Save encoder outputs for source sentences
|
||||
python $SAVE_ENCODER \
|
||||
${INPUT_DIR} \
|
||||
--path ${MODEL} \
|
||||
--task translation_multi_simple_epoch \
|
||||
--lang-pairs ${source_lang}-${target_lang} \
|
||||
--lang-dict ${LANG_DICT} \
|
||||
--gen-subset ${SPLIT} \
|
||||
--bpe 'sentencepiece' \
|
||||
-s ${source_lang} -t ${target_lang} \
|
||||
--sentencepiece-model ${SPM} \
|
||||
--remove-bpe 'sentencepiece' \
|
||||
--beam 1 \
|
||||
--lang-tok-style mbart \
|
||||
--encoder-save-dir ${ENCODER_SAVE_DIR}/${source_lang}
|
||||
|
||||
## Save encoder outputs for target sentences
|
||||
python $SAVE_ENCODER \
|
||||
${INPUT_DIR} \
|
||||
--path ${MODEL} \
|
||||
--lang-pairs ${source_lang}-${target_lang} \
|
||||
--lang-dict ${LANG_DICT} \
|
||||
--task translation_multi_simple_epoch \
|
||||
--gen-subset ${SPLIT} \
|
||||
--bpe 'sentencepiece' \
|
||||
-t ${source_lang} -s ${target_lang} \
|
||||
--sentencepiece-model ${SPM} \
|
||||
--remove-bpe 'sentencepiece' \
|
||||
--beam 1 \
|
||||
--lang-tok-style mbart \
|
||||
--encoder-save-dir ${ENCODER_SAVE_DIR}/${target_lang}
|
||||
|
||||
## Mining
|
||||
python mining/mine.py \
|
||||
--src-lang ${source_lang} \
|
||||
--tgt-lang ${target_lang} \
|
||||
--dim 1024 \
|
||||
--mem 10 \
|
||||
--neighborhood 4 \
|
||||
--src-dir ${ENCODER_SAVE_DIR}/${source_lang} \
|
||||
--tgt-dir ${ENCODER_SAVE_DIR}/${target_lang} \
|
||||
--output $SAVE_DIR \
|
||||
--threshold ${THRESHOLD} \
|
||||
--min-count ${MIN_COUNT} \
|
||||
--valid-size 100 \
|
||||
--dict-path ${DICT} \
|
||||
--spm-path ${SPM} \
|
||||
|
||||
|
||||
## Process and binarize mined data
|
||||
python $SPM_ENCODE \
|
||||
--model ${SPM} \
|
||||
--output_format=piece \
|
||||
--inputs mining/${source_lang}_${target_lang}_mined/train.${source_lang} mining/${source_lang}_${target_lang}_mined/train.${target_lang} \
|
||||
--outputs mining/${source_lang}_${target_lang}_mined/train.bpe.${source_lang} mining/${source_lang}_${target_lang}_mined/train.bpe.${target_lang}
|
||||
|
||||
python $SPM_ENCODE \
|
||||
--model ${SPM} \
|
||||
--output_format=piece \
|
||||
--inputs mining/${source_lang}_${target_lang}_mined/valid.${source_lang} mining/${source_lang}_${target_lang}_mined/valid.${target_lang} \
|
||||
--outputs mining/${source_lang}_${target_lang}_mined/valid.bpe.${source_lang} mining/${source_lang}_${target_lang}_mined/valid.bpe.${target_lang}
|
||||
|
||||
|
||||
fairseq-preprocess \
|
||||
--source-lang ${source_lang} \
|
||||
--target-lang ${target_lang} \
|
||||
--trainpref mining/${source_lang}_${target_lang}_mined/train.bpe \
|
||||
--validpref mining/${source_lang}_${target_lang}_mined/valid.bpe \
|
||||
--destdir mining/${source_lang}_${target_lang}_mined \
|
||||
--srcdict ${DICT} \
|
||||
--joined-dictionary \
|
||||
--workers 8
|
||||
Reference in New Issue
Block a user