test

run_flax_speech_recognition_seq2seq_streaming.py

@@ -716,35 +716,33 @@ def main():
         return {"wer": wer, "cer": cer}
 
     # TODO: Sanchit. This part is not essential and a bit verbose.
-    # 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_samples,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
-        inference_df = pd.DataFrame(columns=['wer','target','prediction'])
+        # Put all predictions into a table
+        inference_df = pd.DataFrame(columns=['mp3','target','prediction'])
         
         idx = 0
         for pred,label in zip(pred_ids,label_ids):
             pred_text = tokenizer.decode(pred,skip_special_tokens=True)
             label_text = tokenizer.decode(label,skip_special_tokens=True)
-            wer = 100 * metric_wer.compute(predictions=[pred_text], references=[label_text])
-            wer_formatted = f'{wer:.2f}'
-            audio_control = f'<audio controls><source src="{os.path.join(predictions_folder_name, "mp3")+"/pred_"+str(idx)+".mp3"}" type="audio/mpeg"></audio>'
-            new_row = pd.DataFrame({'mp3': audio_control, 'wer': wer_formatted, 'target': label_text, 'prediction': pred_text}, index=[0])
+            #wer = 100 * metric_wer.compute(predictions=[pred_text], references=[label_text])
+            #wer_formatted = f'{wer:.2f}'
+            audio_control = f'[↓]({os.path.join(predictions_folder_name, "mp3")+"/pred_"+str(idx)+".mp3"})'
+            new_row = pd.DataFrame({'mp3': audio_control, 'target': label_text, 'prediction': pred_text}, index=[0])
             inference_df = pd.concat([inference_df, new_row], ignore_index=True)
             idx += 1
         
         # Create the prediction table of the first N rows
-        inference_df = inference_df[['mp3', 'wer', 'target','prediction']]
+        inference_df = inference_df[['mp3', 'target','prediction']]
         predict_table = inference_df[0:data_args.number_write_predictions].to_markdown(index=False)
         
         # Build the markdown page
         markdown_str = f"{eval_table}\n\n{predict_table}"
         
-        
         # Save the stats file
         stats_file_name = f"{predictions_folder_name}/step_{step}.md"
         with open(stats_file_name, "w") as f:
@@ -1042,7 +1040,7 @@ def main():
                     state.params, batch.data, min_device_batch=training_args.per_device_eval_batch_size
                 )
                 eval_metrics.append(metrics)
-                if training_args.predict_with_generate and data_args.number_write_predictions:
+                if training_args.predict_with_generate and data_args.number_write_predictions and len(eval_samples)<data_args.number_write_predictions+eval_batch_size:
                     eval_samples.extend(samples['input_features'])
                                 
                 # generation