move creation of linregress.json to test step

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]: ...

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
 
@@ -136,7 +136,7 @@ def translate(self):
                     self.produce_multiple_translations,
                     self.save_confidences,
                     bool(quality_estimation),
-                    verse_test_scores_path,
+                    linregress_path,
                     postprocess_handler,
                     translate_config.get("tags"),
                 )
@@ -149,7 +149,7 @@ def translate(self):
                     self.produce_multiple_translations,
                     self.save_confidences,
                     bool(quality_estimation),
-                    verse_test_scores_path,
+                    linregress_path,
                     postprocess_handler,
                     translate_config.get("tags"),
                 )
@@ -198,7 +198,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",

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_FILENAME
 
 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,52 +156,40 @@ 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():
+        LOGGER.info(f"Searching for {LINREGRESS_FILENAME} in directory {linregress_path}.")
+        linregress_file = linregress_path / LINREGRESS_FILENAME
+        if not linregress_file.is_file():
+            raise ValueError(f"No {LINREGRESS_FILENAME} file found in directory {linregress_path}.")
+        linregress_path = linregress_file
+        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.")
     if not all(cf.is_file() for cf in confidence_file_paths):
         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 +219,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 +410,11 @@ 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.json file containing the confidence-to-chrF3 line of "
+        + "best fit produced by the test step, e.g., project_folder/exp_folder/linregress.json. "
+        + "If a directory is provided instead, linregress.json within it is used.",
     )
     parser.add_argument(
         "confidence_files",
@@ -523,7 +483,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__":

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_FILENAME = "linregress.json"
 
 
 class PairScore:
@@ -280,6 +282,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 +327,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 +353,17 @@ def write_pair_verse_scores(
                 row.append(sentence.rstrip("\n"))
             writer.writerow(row)
 
+    if compute_linregress:
+        write_linregress(linregress_chrf3_scores, linregress_confidence_scores, config.exp_dir / LINREGRESS_FILENAME)
+
+
+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 score_individual_books(
     book_dict: Dict[str, Tuple[List[str], List[List[str]], List[float]]],

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
@@ -99,7 +99,7 @@ def translate_books(
         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,
@@ -180,9 +180,14 @@ def translate_books(
                 raise RuntimeError(f"Some books failed to translate: {' '.join(translation_failed)}")
 
         if quality_estimation and len(confidence_files) > 0:
-            LOGGER.info("Running quality estimation...")
-            estimate_quality(verse_test_scores_path, confidence_files)
-            LOGGER.info("Quality estimation completed.")
+            if linregress_path is None:
+                LOGGER.warning(
+                    "Skipping quality estimation because no linear regression file was provided."
+                )
+            else:
+                LOGGER.info("Running quality estimation...")
+                estimate_quality(linregress_path, confidence_files)
+                LOGGER.info("Quality estimation completed.")
 
     def translate_files(
         self,
@@ -193,7 +198,7 @@ def translate_files(
         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,
@@ -288,9 +293,14 @@ def translate_files(
                     confidence_files.extend(trg_file_path.parent.glob(f"{trg_file_path.stem}*{CONFIDENCE_SUFFIX}"))
 
         if quality_estimation and len(confidence_files) > 0:
-            LOGGER.info("Running quality estimation...")
-            estimate_quality(verse_test_scores_path, confidence_files)
-            LOGGER.info("Quality estimation completed.")
+            if linregress_path is None:
+                LOGGER.warning(
+                    "Skipping quality estimation because no linear regression file was provided."
+                )
+            else:
+                LOGGER.info("Running quality estimation...")
+                estimate_quality(linregress_path, confidence_files)
+                LOGGER.info("Quality estimation completed.")
 
     def _init_translation_task(self, experiment_suffix: str) -> Tuple[Translator, Config, str]:
         clearml = SILClearML(
@@ -438,12 +448,12 @@ 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="The 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.json`. "
+        + "If not provided, the experiment directory will be used to locate linregress.json.",
     )
     parser.add_argument(
         "--debug",
@@ -462,14 +472,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()
 
@@ -501,7 +511,7 @@ def main() -> None:
             args.multiple_translations,
             args.save_confidences,
             args.quality_estimation,
-            verse_test_scores_path,
+            linregress_path,
             postprocess_handler,
             vref=args.vref,
         )
@@ -521,7 +531,7 @@ def main() -> None:
             args.multiple_translations,
             args.save_confidences,
             args.quality_estimation,
-            verse_test_scores_path,
+            linregress_path,
             postprocess_handler,
             vref=args.vref,
         )