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
 
@@ -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",

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,52 +156,41 @@ 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.")
     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 +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__":

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,
             )
         )