| """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() | |