Update app.py

app.py

@@ -1,5 +1,6 @@
-import torch
+import re
 
+import torch
 import gradio as gr
 import yt_dlp as youtube_dl
 from transformers import pipeline
@@ -42,11 +43,8 @@ def transcribe(inputs, prompt):
     prompt = "。" if not prompt else prompt
     generate_kwargs['prompt_ids'] = pipe.tokenizer.get_prompt_ids(prompt, return_tensors='pt').to(device)
     text = pipe(inputs, generate_kwargs=generate_kwargs)['text']
-    if text.startswith(f" {prompt}"):
-        text = text[1 + len(prompt):]
-    elif text.startswith(prompt"):
-        text = text[len(prompt):]
-    return text
+    # currently the pipeline for ASR appends the prompt at the beginning of the transcription, so remove it
+    return re.sub(rf"\A\s*{prompt}\s*", "", text)
 
 def _return_yt_html_embed(yt_url):
     video_id = yt_url.split("?v=")[-1]
@@ -91,11 +89,8 @@ def yt_transcribe(yt_url, prompt, max_filesize=75.0):
     prompt = "。" if not prompt else prompt
     generate_kwargs['prompt_ids'] = pipe.tokenizer.get_prompt_ids(prompt, return_tensors='pt').to(device)
     text = pipe(inputs, generate_kwargs=generate_kwargs)['text']
-    if text.startswith(f" {prompt}"):
-        text = text[1 + len(prompt):]
-    elif text.startswith(prompt"):
-        text = text[len(prompt):]
-    return html_embed_str, text
+    # currently the pipeline for ASR appends the prompt at the beginning of the transcription, so remove it
+    return html_embed_str, re.sub(rf"\A\s*{prompt}\s*", "", text)
 
 
 demo = gr.Blocks()