Update app.py

app.py

@@ -41,8 +41,12 @@ def transcribe(inputs, prompt):
     generate_kwargs = {"language": "japanese", "task": "transcribe"}
     prompt = "。" if not prompt else prompt
     generate_kwargs['prompt_ids'] = pipe.tokenizer.get_prompt_ids(prompt, return_tensors='pt').to(device)
-    result = pipe(inputs, generate_kwargs=generate_kwargs)
-    return result['text'][1 + len(prompt) + 1:] 
+    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
 
 def _return_yt_html_embed(yt_url):
     video_id = yt_url.split("?v=")[-1]
@@ -86,8 +90,12 @@ def yt_transcribe(yt_url, prompt, max_filesize=75.0):
     generate_kwargs = {"language": "japanese", "task": "transcribe"}
     prompt = "。" if not prompt else prompt
     generate_kwargs['prompt_ids'] = pipe.tokenizer.get_prompt_ids(prompt, return_tensors='pt').to(device)
-    result = pipe(inputs, generate_kwargs=generate_kwargs)
-    return html_embed_str, result['text'][1 + len(prompt) + 1:] 
+    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
 
 
 demo = gr.Blocks()