diff --git a/silnlp/common/linear_regression.py b/silnlp/common/linear_regression.py index 596ab38a..f1f2290a 100644 --- a/silnlp/common/linear_regression.py +++ b/silnlp/common/linear_regression.py @@ -19,6 +19,14 @@ class LinearRegressionResult: def toJSON(self) -> str: return json.dumps({"version": self.version, "slope": self.slope, "intercept": self.intercept}, indent=2) + @classmethod + def fromJSON(cls, json_str: str) -> "LinearRegressionResult": + data = json.loads(json_str) + try: + return cls(version=data["version"], slope=float(data["slope"]), intercept=float(data["intercept"])) + except (KeyError, TypeError, ValueError) as e: + raise ValueError(f"Invalid linear regression data: {json_str}") from e + class PointWeightingScheme: def weight_points(self, _x: List[float], _y: List[float]) -> List[float]: ... diff --git a/silnlp/nmt/experiment.py b/silnlp/nmt/experiment.py index 608fb4df..b7a627dd 100644 --- a/silnlp/nmt/experiment.py +++ b/silnlp/nmt/experiment.py @@ -99,13 +99,13 @@ def translate(self): quality_estimation = translate_configs.get("quality_estimation", False) if quality_estimation: - verse_test_scores_path = self.environment.get_mt_exp_dir( - quality_estimation.get("verse_test_scores_file") - if isinstance(quality_estimation, dict) and quality_estimation.get("verse_test_scores_file") - else self.config.exp_dir + linregress_path = self.environment.get_mt_exp_dir( + str(quality_estimation.get("linregress_file")) + if isinstance(quality_estimation, dict) and quality_estimation.get("linregress_file") + else str(self.config.exp_dir) ) else: - verse_test_scores_path = None + linregress_path = None if quality_estimation and not self.save_confidences: self.save_confidences = True @@ -135,8 +135,7 @@ def translate(self): translate_config.get("trg_iso"), self.produce_multiple_translations, self.save_confidences, - bool(quality_estimation), - verse_test_scores_path, + linregress_path, postprocess_handler, translate_config.get("tags"), ) @@ -148,8 +147,7 @@ def translate(self): translate_config.get("trg_iso"), self.produce_multiple_translations, self.save_confidences, - bool(quality_estimation), - verse_test_scores_path, + linregress_path, postprocess_handler, translate_config.get("tags"), ) @@ -198,7 +196,8 @@ def main() -> None: "--multiple-translations", default=False, action="store_true", - help='Produce multiple translations of each verse. These will be saved in separate files with suffixes like ".1.txt", ".2.txt", etc.', + help='Produce multiple translations of each verse. ' + + 'These will be saved in separate files with suffixes like ".1.txt", ".2.txt", etc.', ) parser.add_argument( "--save-confidences", diff --git a/silnlp/nmt/quality_estimation.py b/silnlp/nmt/quality_estimation.py index abb2cc5e..6f3ad693 100644 --- a/silnlp/nmt/quality_estimation.py +++ b/silnlp/nmt/quality_estimation.py @@ -10,9 +10,9 @@ from machine.scripture import ALL_BOOK_IDS, VerseRef from ..common.environment import SilNlpEnv -from ..common.linear_regression import perform_enhanced_linear_regression +from ..common.linear_regression import LinearRegressionResult from ..common.translator import CONFIDENCE_SUFFIX, ConfidenceFile, TxtConfidenceFile, UsfmConfidenceFile -from .test import VERSE_SCORES_SUFFIX +from .test import LINREGRESS_PREFIX LOGGER = logging.getLogger(__package__ + ".quality_estimation") CANONICAL_ORDER = {book: i for i, book in enumerate(ALL_BOOK_IDS)} @@ -95,7 +95,7 @@ def add_scores_from_confidence_file( @dataclass class SequenceScore(Score): - sequence_num: str + sequence_num: int trg_draft_file_stem: str @classmethod @@ -140,10 +140,10 @@ def add_scores_from_confidence_file( self.add_score(trg_draft_file_stem, score) -def estimate_quality(verse_test_scores_path: Path, confidence_file_paths: List[Path]) -> None: - verse_test_scores_path, confidence_files = validate_inputs(verse_test_scores_path, confidence_file_paths) +def estimate_quality(linregress_path: Path, confidence_file_paths: List[Path]) -> None: + linear_regression_result, confidence_files = validate_inputs(linregress_path, confidence_file_paths) verse_scores, chapter_scores, book_scores, sequence_scores, txt_file_scores = project_chrf3( - verse_test_scores_path, confidence_files + linear_regression_result, confidence_files ) compute_usable_proportions( verse_scores, @@ -156,19 +156,18 @@ def estimate_quality(verse_test_scores_path: Path, confidence_file_paths: List[P def validate_inputs( - verse_test_scores_path: Path, confidence_file_paths: List[Path] -) -> Tuple[Path, List[ConfidenceFile]]: - if not verse_test_scores_path.exists(): - raise FileNotFoundError(f"Test data file {verse_test_scores_path} does not exist.") - elif verse_test_scores_path.is_dir(): - LOGGER.info(f"Searching for files with suffix {VERSE_SCORES_SUFFIX} in directory {verse_test_scores_path}.") - test_files = list(verse_test_scores_path.glob(f"*{VERSE_SCORES_SUFFIX}")) - if not test_files: - raise ValueError( - f"No test data file with the {VERSE_SCORES_SUFFIX} suffix found in directory {verse_test_scores_path}." - ) - verse_test_scores_path = test_files[0] - LOGGER.info(f"Using test data file {verse_test_scores_path}.") + linregress_path: Path, confidence_file_paths: List[Path] +) -> Tuple[LinearRegressionResult, List[ConfidenceFile]]: + if not linregress_path.exists(): + raise FileNotFoundError(f"Linear regression file {linregress_path} does not exist.") + elif linregress_path.is_dir(): + pattern = f"{LINREGRESS_PREFIX}.*.json" + LOGGER.info(f"Searching for files matching {pattern} in directory {linregress_path}.") + linregress_files = list(linregress_path.glob(pattern)) + if not linregress_files: + raise ValueError(f"No file matching {pattern} found in directory {linregress_path}.") + linregress_path = linregress_files[0] + LOGGER.info(f"Using linear regression file {linregress_path}.") if len(confidence_file_paths) == 0: raise ValueError("At least one confidence file must be provided.") @@ -176,32 +175,22 @@ def validate_inputs( missing_files = [str(cf) for cf in confidence_file_paths if not cf.is_file()] raise FileNotFoundError(f"The following confidence files do not exist: {', '.join(missing_files)}") + with open(linregress_path, "r", encoding="utf-8") as f: + linear_regression_result = LinearRegressionResult.fromJSON(f.read()) + confidence_files: List[ConfidenceFile] = [] for cf in confidence_file_paths: confidence_files.append(ConfidenceFile.from_confidence_file_path(cf)) - return verse_test_scores_path, confidence_files + return linear_regression_result, confidence_files def project_chrf3( - verse_test_scores_path: Path, confidence_files: List[ConfidenceFile] + linear_regression_result: LinearRegressionResult, confidence_files: List[ConfidenceFile] ) -> Tuple[List[VerseScore], ChapterScores, BookScores, List[SequenceScore], TxtFileScores]: - chrf3_scores, confidence_scores = extract_test_data(verse_test_scores_path) - if len(chrf3_scores) != len(confidence_scores): - raise ValueError( - f"The number of chrF3 scores ({len(chrf3_scores)}) and confidence scores ({len(confidence_scores)}) " - f"in {verse_test_scores_path} do not match." - ) - - linear_regression_result = perform_enhanced_linear_regression(confidence_scores, chrf3_scores) slope = linear_regression_result.slope intercept = linear_regression_result.intercept LOGGER.info(f"Linear regression data:\n{linear_regression_result.toJSON()}") - output_dir = confidence_files[0].get_path().parent - output_file = output_dir / "linregress.json" - with open(output_file, "w", encoding="utf-8") as f: - LOGGER.info(f"Saving linear regression data to {output_file}") - f.write(linear_regression_result.toJSON()) verse_scores: List[VerseScore] = [] chapter_scores: ChapterScores = ChapterScores() @@ -231,34 +220,6 @@ def project_chrf3( return verse_scores, chapter_scores, book_scores, sequence_scores, txt_file_scores -def extract_test_data(verse_test_scores_path: Path) -> Tuple[List[float], List[float]]: - chrf3_scores: List[float] = [] - confidence_scores: List[float] = [] - with open(verse_test_scores_path, "r", encoding="utf-8") as f: - header = next(f).strip().lower().split("\t") - try: - chrf3_index = header.index("chrf3") - confidence_index = header.index("confidence") - except ValueError as e: - raise ValueError( - f"Could not find 'chrF3' and/or 'confidence' columns in header of {verse_test_scores_path}: {header}" - ) from e - for line_num, line in enumerate(f, start=2): - cols = line.strip().split("\t") - try: - chrf3 = float(cols[chrf3_index]) - confidence = float(cols[confidence_index]) - chrf3_scores.append(chrf3) - confidence_scores.append(confidence) - except (ValueError, IndexError) as e: - raise ValueError( - f"Error parsing line {line_num} in {verse_test_scores_path}: {line.strip()}" - f" (chrF3 index: {chrf3_index}, confidence index: {confidence_index})" - ) from e - - return chrf3_scores, confidence_scores - - @dataclass class UsabilityParameters: count: float @@ -450,11 +411,12 @@ def compute_txt_file_usability( def main() -> None: parser = argparse.ArgumentParser(description="Estimate the quality of drafts created by an NMT model.") parser.add_argument( - "verse_test_scores_file", + "linregress_file", type=str, - help="Path relative to MT/experiments to a verse-level test score file, which is used to find line of best fit " - + "for confidence and chrF3, e.g., project_folder/exp_folder/test.trg-predictions.detok.txt.5000.scores.tsv. " - + "If a directory is provided instead, the first *.scores.tsv match is used.", + help="Path relative to MT/experiments to a linregress file containing the confidence-to-chrF3 line of best " + + f"fit produced by the test step, e.g., project_folder/exp_folder/{LINREGRESS_PREFIX}.5000.json (or " + + f"{LINREGRESS_PREFIX}.eng.fra.5000.json for an experiment with multiple language pairs). " + + f"If a directory is provided instead, the first {LINREGRESS_PREFIX}.*.json match is used.", ) parser.add_argument( "confidence_files", @@ -523,7 +485,7 @@ def main() -> None: for book_id in args.books: confidence_file_paths.extend(confidence_dir.glob(f"[0-9]*{book_id}{draft_suffix}.*{CONFIDENCE_SUFFIX}")) - estimate_quality(environment.get_mt_exp_dir(args.verse_test_scores_file), confidence_file_paths) + estimate_quality(environment.get_mt_exp_dir(args.linregress_file), confidence_file_paths) if __name__ == "__main__": diff --git a/silnlp/nmt/test.py b/silnlp/nmt/test.py index 91999213..8f7cb473 100644 --- a/silnlp/nmt/test.py +++ b/silnlp/nmt/test.py @@ -13,6 +13,7 @@ from scipy.stats import gmean from ..common.environment import SilNlpEnv +from ..common.linear_regression import perform_enhanced_linear_regression from ..common.translator import CONFIDENCE_SUFFIX from ..common.utils import get_git_revision_hash from .clearml_connection import TAGS_LIST, SILClearML @@ -50,6 +51,7 @@ TEST_TRG_PREDICTIONS_PREFIX = "test.trg-predictions" VERSE_SCORES_SUFFIX = ".scores.tsv" +LINREGRESS_PREFIX = "linregress" class PairScore: @@ -120,6 +122,7 @@ def score_pair( ref_projects: Set[str], draft_index: int = 1, pair_confs: Optional[List[float]] = None, + linregress_file_name: Optional[str] = None, ) -> PairScore: bleu_score = None if "bleu" in scorers: @@ -242,6 +245,7 @@ def score_pair( other_scores, config, confidences if "confidence" in scorers else None, + linregress_file_name, ) return PairScore(book, src_iso, trg_iso, bleu_score, len(pair_sys), ref_projects, other_scores, draft_index) @@ -256,6 +260,7 @@ def write_pair_verse_scores( other_scores: Dict[str, float], config: Config, confidences: Optional[List[float]], + linregress_file_name: Optional[str] = None, ) -> None: scorers = scorers.intersection(SUPPORTED_SENTENCE_SCORERS) other_scores = {k: v for k, v in other_scores.items() if k.lower() in scorers} @@ -280,6 +285,9 @@ def write_pair_verse_scores( header.append("Reference") writer.writerow(header) spbleu_metric = sacrebleu.metrics.BLEU(tokenize="flores200", lowercase=True) if "spbleu" in scorers else None + compute_linregress = "chrf3" in scorers and "confidence" in scorers and confidences is not None + linregress_chrf3_scores: List[float] = [] + linregress_confidence_scores: List[float] = [] for index, pred in enumerate(pair_sys): sentences: List[str] = [] for ref in pair_refs: @@ -322,6 +330,10 @@ def write_pair_verse_scores( if "confidence" in scorers and confidences is not None: other_verse_scores["Confidence"] = confidences[index] + if compute_linregress: + linregress_chrf3_scores.append(other_verse_scores["chrF3"]) + linregress_confidence_scores.append(other_verse_scores["Confidence"]) + row: List[str] = [f"{index + 1}"] if "bleu" in scorers: @@ -344,6 +356,38 @@ def write_pair_verse_scores( row.append(sentence.rstrip("\n")) writer.writerow(row) + if compute_linregress and linregress_file_name is not None: + write_linregress( + linregress_chrf3_scores, + linregress_confidence_scores, + config.exp_dir / linregress_file_name, + ) + + +def write_linregress(chrf3_scores: List[float], confidence_scores: List[float], output_path: Path) -> None: + linear_regression_result = perform_enhanced_linear_regression(confidence_scores, chrf3_scores) + LOGGER.info(f"Linear regression data:\n{linear_regression_result.toJSON()}") + LOGGER.info(f"Saving linear regression data to {output_path}") + with open(output_path, "w", encoding="utf-8") as f: + f.write(linear_regression_result.toJSON()) + + +def get_linregress_file_name( + step_token: str, + split_by_pair: bool, + src_iso: str, + trg_iso: str, + produce_multiple_translations: bool, + draft_index: int, +) -> str: + parts = [LINREGRESS_PREFIX] + if split_by_pair: + parts.extend([src_iso, trg_iso]) + parts.append(step_token) + if produce_multiple_translations: + parts.append(str(draft_index)) + return ".".join(parts) + ".json" + def score_individual_books( book_dict: Dict[str, Tuple[List[str], List[List[str]], List[float]]], @@ -538,10 +582,11 @@ def test_checkpoint( refs_patterns: List[str] = [] translation_detok_file_names: List[str] = [] translation_conf_file_names: List[str] = [] + step_token = "avg" if step == -1 else str(step) suffix_str = "_".join(map(lambda n: book_number_to_id(n), sorted(books.keys()))) if len(suffix_str) > 0: suffix_str += "-" - suffix_str += "avg" if step == -1 else str(step) + suffix_str += step_token features_file_name = "test.src.txt" if (config.exp_dir / features_file_name).is_file(): @@ -633,11 +678,16 @@ def test_checkpoint( ): src_iso = config.default_test_src_iso trg_iso = config.default_test_trg_iso - if features_file_name != "test.src.txt": + split_by_pair = features_file_name != "test.src.txt" + if split_by_pair: parts = features_file_name.split(".") src_iso = parts[1] trg_iso = parts[2] + linregress_file_name = get_linregress_file_name( + step_token, split_by_pair, src_iso, trg_iso, produce_multiple_translations, draft_index + ) + pair_sys, pair_refs, book_dict = load_test_data( tokenizer, vref_file_name, @@ -674,6 +724,7 @@ def test_checkpoint( config, ref_projects, draft_index, + linregress_file_name=linregress_file_name, ) ) diff --git a/silnlp/nmt/translate.py b/silnlp/nmt/translate.py index 8304dd45..80aaa56d 100644 --- a/silnlp/nmt/translate.py +++ b/silnlp/nmt/translate.py @@ -6,7 +6,7 @@ from typing import Generator, Iterable, List, Optional, Tuple, Union from machine.corpora import UsfmFileTextCorpus, create_versification_ref_corpus, extract_scripture_corpus -from machine.scripture import VerseRef, book_number_to_id, get_chapters +from machine.scripture import book_number_to_id, get_chapters from ..common.environment import SilNlpEnv from ..common.paratext import book_file_name_digits @@ -98,8 +98,7 @@ def translate_books( trg_iso: Optional[str], produce_multiple_translations: bool = False, save_confidences: bool = False, - quality_estimation: bool = False, - verse_test_scores_path: Optional[Path] = None, + linregress_path: Optional[Path] = None, postprocess_handler: Optional[PostprocessHandler] = None, tags: Optional[List[str]] = None, vref: bool = False, @@ -179,9 +178,9 @@ def translate_books( if len(translation_failed) > 0: raise RuntimeError(f"Some books failed to translate: {' '.join(translation_failed)}") - if quality_estimation and len(confidence_files) > 0: + if linregress_path is not None: LOGGER.info("Running quality estimation...") - estimate_quality(verse_test_scores_path, confidence_files) + estimate_quality(linregress_path, confidence_files) LOGGER.info("Quality estimation completed.") def translate_files( @@ -192,8 +191,7 @@ def translate_files( trg_iso: Optional[str], produce_multiple_translations: bool = False, save_confidences: bool = False, - quality_estimation: bool = False, - verse_test_scores_path: Optional[Path] = None, + linregress_path: Optional[Path] = None, postprocess_handler: Optional[PostprocessHandler] = None, tags: Optional[List[str]] = None, vref: bool = False, @@ -287,9 +285,9 @@ def translate_files( if save_confidences: confidence_files.extend(trg_file_path.parent.glob(f"{trg_file_path.stem}*{CONFIDENCE_SUFFIX}")) - if quality_estimation and len(confidence_files) > 0: + if linregress_path is not None: LOGGER.info("Running quality estimation...") - estimate_quality(verse_test_scores_path, confidence_files) + estimate_quality(linregress_path, confidence_files) LOGGER.info("Quality estimation completed.") def _init_translation_task(self, experiment_suffix: str) -> Tuple[Translator, Config, str]: @@ -438,12 +436,13 @@ def main() -> None: help="Run quality estimation after translation completes. Requires --save-confidences.", ) parser.add_argument( - "--verse-test-scores-file", + "--linregress-file", type=str, default=None, - help="The tsv file relative to MT/experiments containing the verse-level test scores to determine " - + "line of best fit, e.g., `project_folder/exp_folder/test.trg-predictions.detok.txt.5000.scores.tsv`. " - + "If not provided, the experiment directory will be used to locate the test scores file.", + help="A linregress.*.json file relative to MT/experiments containing the confidence-to-chrF3 line of best " + + "fit coefficients produced by the test step, e.g., `project_folder/exp_folder/linregress.5000.json` (or " + + "`linregress.eng.fra.5000.json` for an experiment with multiple language pairs). " + + "If not provided, the experiment directory will be searched for a linregress.*.json file.", ) parser.add_argument( "--debug", @@ -462,14 +461,14 @@ def main() -> None: if not args.save_confidences: args.save_confidences = True - if args.verse_test_scores_file is None: - verse_test_scores_path = environment.get_mt_exp_dir(args.experiment) + if args.linregress_file is None: + linregress_path = environment.get_mt_exp_dir(args.experiment) else: - verse_test_scores_path = environment.get_mt_exp_dir(args.verse_test_scores_file) - if not verse_test_scores_path.exists(): - parser.error(f"The verse test scores path {verse_test_scores_path} does not exist.") + linregress_path = environment.get_mt_exp_dir(args.linregress_file) + if not linregress_path.exists(): + parser.error(f"The linear regression path {linregress_path} does not exist.") else: - verse_test_scores_path = None + linregress_path = None get_git_revision_hash() @@ -500,8 +499,7 @@ def main() -> None: args.trg_iso, args.multiple_translations, args.save_confidences, - args.quality_estimation, - verse_test_scores_path, + linregress_path, postprocess_handler, vref=args.vref, ) @@ -520,8 +518,7 @@ def main() -> None: args.trg_iso, args.multiple_translations, args.save_confidences, - args.quality_estimation, - verse_test_scores_path, + linregress_path, postprocess_handler, vref=args.vref, )