update readme

README.md

@@ -41,6 +41,7 @@ from experiments.utils import set_logger, get_device, remove_suppress_tokens
 from experiments.utils.utils import UNSUPPRESS_TOKEN
 import torchaudio
 import numpy as np
+
 set_logger()
 
 
@@ -63,7 +64,9 @@ def main(model_path, audio_file_path, prompt, max_new_tokens, language, device):
     if signal.ndim == 2:
         signal = torch.mean(signal, dim=0)
     # pre-process to get the input features
-    input_features = processor(signal, sampling_rate=target_sample_rate, return_tensors="pt").input_features
+    input_features = processor(
+        signal, sampling_rate=target_sample_rate, return_tensors="pt"
+    ).input_features
     input_features = input_features.to(device)
 
     prompt = prompt.lower()  # lowercase the prompt, to align with training
@@ -74,8 +77,11 @@ def main(model_path, audio_file_path, prompt, max_new_tokens, language, device):
     # generate token ids by running model forward sequentially
     logging.info(f"Inference with prompt: '{prompt}'.")
     predicted_ids = model.generate(
-        input_features, max_new_tokens=max_new_tokens, language=language,
-        prompt_ids=prompt_ids, generation_config=model.generation_config
+        input_features,
+        max_new_tokens=max_new_tokens,
+        language=language,
+        prompt_ids=prompt_ids,
+        generation_config=model.generation_config,
     )
 
     # post-process token ids to text
@@ -84,21 +90,50 @@ def main(model_path, audio_file_path, prompt, max_new_tokens, language, device):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Transcribe audio using Whisper model.")
-    parser.add_argument('--model-path', type=str, 
-                        required=True, 
-                        default="aiola/whisper-ner-v1",
-                        help='Path to the pre-trained model components.')
-    parser.add_argument('--audio-file-path',
-                        type=str, 
-                        required=True, 
-                        help='Path to the audio file to transcribecd.')
-    parser.add_argument('--prompt', type=str, default='father', help='Prompt text to guide the transcription.')
-    parser.add_argument('--max-new-tokens', type=int, default=256, help='Maximum number of new tokens to generate.')
-    parser.add_argument('--language', type=str, default='en', help='Language code for the transcription.')
+    parser = argparse.ArgumentParser(
+        description="Transcribe audio using Whisper model."
+    )
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        required=True,
+        default="aiola/whisper-ner-v1",
+        help="Path to the pre-trained model components.",
+    )
+    parser.add_argument(
+        "--audio-file-path",
+        type=str,
+        required=True,
+        help="Path to the audio file (wav) to transcribe.",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="father",
+        help="Prompt text to guide the transcription.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=256,
+        help="Maximum number of new tokens to generate.",
+    )
+    parser.add_argument(
+        "--language",
+        type=str,
+        default="en",
+        help="Language code for the transcription.",
+    )
 
     args = parser.parse_args()
     device = get_device()
-    main(args.model_path, args.audio_file_path, args.prompt, args.max_new_tokens, args.language, device)
+    main(
+        args.model_path,
+        args.audio_file_path,
+        args.prompt,
+        args.max_new_tokens,
+        args.language,
+        device,
+    )
 
 ```