test

run_flax_speech_recognition_seq2seq_streaming.py

@@ -717,8 +717,8 @@ def main():
     # We have also considered adding it to Tensorboard logs instead, but it is not so easily browsable. 
     # We do howver think it is extemely useful info for debugging and improving the model
     
-    def write_predictions(step,eval_metrics, pred_ids, label_ids, model_dir):
-        predictions_folder_name = os.path.join(model_dir,"/predictions")
+    def write_predictions(step,eval_metrics, pred_ids, label_ids):
+        predictions_folder_name = os.path.join(training_args.output_dir,"/predictions")
         eval_table = f"| STEP| loss | wer |cer|\n| ---| --- | --- |--- |\n| **{step}**| {eval_metrics['loss']:.3f} | {eval_metrics['wer']:.3f} |{eval_metrics['cer']:.3f} |"
         
         # Put all predictions into a table and calculate wer
@@ -1067,7 +1067,7 @@ def main():
                              eval_metrics, train_time, step)
                 
             if training_args.predict_with_generate and write_predictions:
-                write_predictions(step,eval_metrics, eval_preds, eval_labels,data_args.output_dir)
+                write_predictions(step,eval_metrics, eval_preds, eval_labels)
 
             # save checkpoint after each epoch and push checkpoint to the hub
             if jax.process_index() == 0: