Update README.md

README.md

@@ -219,6 +219,7 @@ print(result["text"])
 Kotoba-whisper can generate transcription with prompting as below:
 
 ```python
+import re
 import torch
 from transformers import pipeline
 from datasets import load_dataset, Audio
@@ -243,16 +244,17 @@ pipe = pipeline(
 dataset = load_dataset("japanese-asr/ja_asr.reazonspeech_test", split="test")
 
 # --- Without prompt ---
-result = pipe(dataset[10]["audio"], generate_kwargs=generate_kwargs)
-print(result['text'])
+text = pipe(dataset[10]["audio"], generate_kwargs=generate_kwargs)['text']
+print(text)
 # 81歳、力強い走りに変わってきます。
 
 # --- With prompt ---: Let's change `81` to `91`.
 prompt = "91歳"
 generate_kwargs['prompt_ids'] = pipe.tokenizer.get_prompt_ids(prompt, return_tensors="pt").to(device)
-result = pipe(dataset[10]["audio"], generate_kwargs=generate_kwargs)
-result['text'] = result['text'][1 + len(prompt) + 1:]  # prompt has been added at the beginning of the output now, so remove it.
-print(result['text'])
+text = pipe(dataset[10]["audio"], generate_kwargs=generate_kwargs)['text']
+# currently the pipeline for ASR appends the prompt at the beginning of the transcription, so remove it
+text = re.sub(rf"\A\s*{prompt}\s*", "", text)
+print(text)
 # あっぶったでもスルガさん、91歳、力強い走りに変わってきます。
 ```