Fix url problem

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import requests
+import re
 from transformers import AutoTokenizer, pipeline
 from youtube_transcript_api._transcripts import TranscriptListFetcher
 
@@ -12,6 +13,14 @@ tokenizer = AutoTokenizer.from_pretrained("./checkpoint-6000")
 max_size = 512
 classes = [False, True]
 
+pattern = re.compile(
+    r"(?:https?:\/\/)?(?:[0-9A-Z-]+\.)?(?:youtube|youtu|youtube-nocookie)\.(?:com|be)\/(?:watch\?v=|watch\?.+&v=|embed\/|v\/|.+\?v=)?([^&=\n%\?]{11})"
+)
+
+def video_id(url):
+    p = pattern.match(url)
+    return p.group(1) if p else None
+
 def process(obj):
     o = obj["events"]
     new_l = []
@@ -65,8 +74,12 @@ def get_transcript(video_id, session):
     p = process(obj.json())
     return p
 
-def transcript(video_id):
-  return " ".join(l["w"].strip() for l in get_transcript(video_id, requests.Session()))
+def transcript(url):
+  i = video_id(url)
+  if i:
+    return " ".join(l["w"].strip() for l in get_transcript(i, requests.Session()))
+  else:
+    return "ERROR: Failed to load transcript (it the link a valid youtube url?)..."
 
 def inference(transcript):
   tokens = tokenizer(transcript.split(" "))["input_ids"]
@@ -119,9 +132,9 @@ def predict(transcript):
 with gr.Blocks() as demo:
     with gr.Row():
       with gr.Column():
-        inp = gr.Textbox(label="Video ID or URL", placeholder="Video id", lines=1, max_lines=1)
+        inp = gr.Textbox(label="Video URL", placeholder="Video URL", lines=1, max_lines=1)
         btn = gr.Button("Fetch Transcript")
-        gr.Examples(["xsLJZyih3Ac"], [inp])
+        gr.Examples(["youtu.be/xsLJZyih3Ac"], [inp])
         text = gr.Textbox(label="Transcript", placeholder="<generated transcript>")
         btn.click(fn=transcript, inputs=inp, outputs=text)
       with gr.Column():