wehub-resource-sync
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
#
|
|
# 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.
|
|
|
|
"""Extracts random constraints from reference files."""
|
|
|
|
import argparse
|
|
import random
|
|
import sys
|
|
|
|
from sacrebleu import extract_ngrams
|
|
|
|
|
|
def get_phrase(words, index, length):
|
|
assert index < len(words) - length + 1
|
|
phr = " ".join(words[index : index + length])
|
|
for i in range(index, index + length):
|
|
words.pop(index)
|
|
return phr
|
|
|
|
|
|
def main(args):
|
|
|
|
if args.seed:
|
|
random.seed(args.seed)
|
|
|
|
for line in sys.stdin:
|
|
constraints = []
|
|
|
|
def add_constraint(constraint):
|
|
constraints.append(constraint)
|
|
|
|
source = line.rstrip()
|
|
if "\t" in line:
|
|
source, target = line.split("\t")
|
|
if args.add_sos:
|
|
target = f"<s> {target}"
|
|
if args.add_eos:
|
|
target = f"{target} </s>"
|
|
|
|
if len(target.split()) >= args.len:
|
|
words = [target]
|
|
|
|
num = args.number
|
|
|
|
choices = {}
|
|
for i in range(num):
|
|
if len(words) == 0:
|
|
break
|
|
segmentno = random.choice(range(len(words)))
|
|
segment = words.pop(segmentno)
|
|
tokens = segment.split()
|
|
phrase_index = random.choice(range(len(tokens)))
|
|
choice = " ".join(
|
|
tokens[phrase_index : min(len(tokens), phrase_index + args.len)]
|
|
)
|
|
for j in range(
|
|
phrase_index, min(len(tokens), phrase_index + args.len)
|
|
):
|
|
tokens.pop(phrase_index)
|
|
if phrase_index > 0:
|
|
words.append(" ".join(tokens[0:phrase_index]))
|
|
if phrase_index + 1 < len(tokens):
|
|
words.append(" ".join(tokens[phrase_index:]))
|
|
choices[target.find(choice)] = choice
|
|
|
|
# mask out with spaces
|
|
target = target.replace(choice, " " * len(choice), 1)
|
|
|
|
for key in sorted(choices.keys()):
|
|
add_constraint(choices[key])
|
|
|
|
print(source, *constraints, sep="\t")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--number", "-n", type=int, default=1, help="number of phrases")
|
|
parser.add_argument("--len", "-l", type=int, default=1, help="phrase length")
|
|
parser.add_argument(
|
|
"--add-sos", default=False, action="store_true", help="add <s> token"
|
|
)
|
|
parser.add_argument(
|
|
"--add-eos", default=False, action="store_true", help="add </s> token"
|
|
)
|
|
parser.add_argument("--seed", "-s", default=0, type=int)
|
|
args = parser.parse_args()
|
|
|
|
main(args)
|