File size: 902 Bytes
ed46c5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | """Retrieval baseline: output the translation of the most similar training transliteration
(Indel similarity on space-free strings, same target language)."""
import json
import sys
from rapidfuzz import fuzz, process
train = [json.loads(l) for l in open(sys.argv[1])]
test = [json.loads(l) for l in open(sys.argv[2])]
out = open(sys.argv[3], "w")
by_lang = {}
for r in train:
by_lang.setdefault(r["lang"], []).append(r)
keys = {lang: [r["src"].replace(" ", "") for r in rows] for lang, rows in by_lang.items()}
for r in test:
q = (r.get(sys.argv[4]) if len(sys.argv) > 4 else r["src"]) or ""
match = process.extractOne(q.replace(" ", ""), keys[r["lang"]], scorer=fuzz.ratio)
rec = dict(r)
rec["pred"] = by_lang[r["lang"]][match[2]]["tgt"] if match else ""
rec["retrieval_score"] = match[1] if match else 0
out.write(json.dumps(rec, ensure_ascii=False) + "\n")
out.close()
|