76 lines
1.8 KiB
Python
76 lines
1.8 KiB
Python
from time import time
|
|
|
|
import datasets
|
|
from fire import Fire
|
|
from sacrebleu import BLEU, CHRF, TER
|
|
|
|
from ..translator import Translator
|
|
|
|
bleu = BLEU()
|
|
chrf = CHRF()
|
|
ter = TER()
|
|
|
|
|
|
def eval(
|
|
model_path: str,
|
|
src_lang: str,
|
|
tgt_lang: str,
|
|
output_file: str,
|
|
beam_size: int = 5,
|
|
intra_threads: int = 1,
|
|
inter_threads: int = 6,
|
|
compute_type="auto",
|
|
device: str = "auto",
|
|
max_batch_size: int = 32,
|
|
max_decoding_length: int = 512,
|
|
):
|
|
t = Translator(
|
|
model_path=model_path,
|
|
device=device,
|
|
compute_type=compute_type,
|
|
inter_threads=inter_threads,
|
|
intra_threads=intra_threads,
|
|
)
|
|
|
|
print("Loading flores-devtest data")
|
|
try:
|
|
flores = datasets.load_dataset(
|
|
"facebook/flores",
|
|
f"{src_lang}-{tgt_lang}", # trust_remote_code=True
|
|
)
|
|
except:
|
|
flores = datasets.load_dataset(
|
|
"facebook/flores",
|
|
f"{tgt_lang}-{src_lang}", # trust_remote_code=True
|
|
)
|
|
|
|
src = []
|
|
ref = []
|
|
for i in flores["devtest"]:
|
|
src.append(i[f"sentence_{src_lang}"])
|
|
ref.append(i[f"sentence_{tgt_lang}"])
|
|
|
|
print("Translating")
|
|
mt = t(src, max_batch_size=max_batch_size, max_decoding_length=max_decoding_length, beam_size=beam_size)
|
|
|
|
# Write results to file, for COMET scoring
|
|
if output_file:
|
|
with open(output_file, "wt") as myfile:
|
|
myfile.write("".join([i.replace("\n", "\t") + "\n" for i in mt]))
|
|
|
|
print("Source sample: ", src[:10])
|
|
print("Reference sample: ", ref[:10])
|
|
print("Translation sample: ", mt[:10])
|
|
|
|
print(bleu.corpus_score(mt, [ref]))
|
|
print(chrf.corpus_score(mt, [ref]))
|
|
print(ter.corpus_score(mt, [ref]))
|
|
|
|
|
|
def main():
|
|
Fire(eval)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|