⚡️ Style Transfer: use all data inside evaluator training set to trai…

…n model not only generation-0
arkhn · Oct 17, 2024 · d3fa9b2 · d3fa9b2
1 parent 3b00e43
commit d3fa9b2
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/lib/style-transfer/style_transfer/rb_gen/steps/score.py b/lib/style-transfer/style_transfer/rb_gen/steps/score.py
@@ -29,14 +29,15 @@ def train_eval_model(
     train_gen_dataset = gen_dataset["train"]
     gen_dataset = gen_dataset["test"]
     train_examples = []
-    train_examples.extend(
-        [
-            InputExample(texts=[pred, ground_text], label=0)
-            for pred, ground_text in zip(
-                train_gen_dataset["generation_0"], train_gen_dataset["ground_texts"]
-            )
-        ]
-    )
+    for num_seq in range(cfg.model.num_generated_sequences):
+        train_examples.extend(
+            [
+                InputExample(texts=[pred, ground_text], label=0)
+                for pred, ground_text in zip(
+                    train_gen_dataset[f"generation_{num_seq}"], train_gen_dataset["ground_texts"]
+                )
+            ]
+        )
     if cfg.score.train.use_ground_truth:
         train_examples.extend(
             [