test smaller flashlight beam

rwth-i6 · Jan 10, 2025 · d35e21e · d35e21e
1 parent d1cf549
commit d35e21e
Showing 1 changed file with 28 additions and 1 deletion.
diff --git a/users/zeyer/experiments/exp2024_04_23_baselines/ctc_recog_ext.py b/users/zeyer/experiments/exp2024_04_23_baselines/ctc_recog_ext.py
@@ -82,6 +82,33 @@ def py():
                 f"{prefix}/recog-beam{beam_size}-lm_{lm_out_name}-lmScale{lm_scale}-priorScale{prior_scale}", res.output
             )
 
+        # Flashlight beam search implementation.
+        for prior_scale, lm_scale in [
+            (0.0, 1.0),
+            # (0.2, 2.0),
+        ]:
+            model = get_ctc_with_lm(
+                ctc_model=ctc_model, prior=prior, prior_scale=prior_scale, language_model=lm, lm_scale=lm_scale
+            )
+            res = recog_model(
+                task=task,
+                model=model,
+                recog_def=model_recog_flashlight,
+                config={
+                    "n_best": 32,
+                    "beam_size": 1024,
+                    "beam_size_token": 16,
+                    "beam_threshold": 14,
+                    "batch_size": 5_000 * ctc_model.definition.batch_size_factor,
+                    "torch_amp": {"dtype": "bfloat16"},
+                },
+                search_rqmt={"cpu": 4, "mem": 30, "time": 24, "gpu_mem": 24},
+            )
+            tk.register_output(
+                f"{prefix}/recog-fl-beamToken16-lm_{lm_out_name}-lmScale{lm_scale}-priorScale{prior_scale}",
+                res.output,
+            )
+
         # Flashlight beam search implementation.
         for prior_scale, lm_scale in [
             (0.0, 1.0),
@@ -105,7 +132,7 @@ def py():
                 search_rqmt={"cpu": 4, "mem": 30, "time": 24, "gpu_mem": 24},
             )
             tk.register_output(
-                f"{prefix}/recog-beam{beam_size}-lm_{lm_out_name}-lmScale{lm_scale}-priorScale{prior_scale}", res.output
+                f"{prefix}/recog-fl-beamToken128-lm_{lm_out_name}-lmScale{lm_scale}-priorScale{prior_scale}", res.output
             )