Add timestamps and title support. Use tiny model

app.py

@@ -5,18 +5,35 @@ import whisper
 from pytube import YouTube
 import gradio as gr
 
-infer_model = whisper.load_model("base")
+infer_model = whisper.load_model("tiny")
 
-def infer(link):
+
+def infer(link: str, add_timestamps: bool) -> str:
     audio_path = download_audio(link)
-    if audio_path is None:
+
+    if not audio_path:
         return "Unable to process request."
 
     result = infer_model.transcribe(audio_path)
-    print(result["text"])
-    return result["text"]
 
-def download_audio(link):
+    title = "Content"
+    try:
+        title = audio_path.split("/")[-1]
+        title = title.split(".")[0]
+    except Exception as e:
+        print(f"Unable to extract title. Exception {e}")
+
+    if not add_timestamps:
+        print(result["text"])
+        return title + "\n" + result["text"]
+
+    result_text = title + "\n"
+    for segment in result["segments"]:
+        result_text += f"{float(segment['start']):.2f}s - {float(segment['end']):.2f}s : {segment['text']}\n"
+    return result_text.strip("\n")
+
+
+def download_audio(link: str) -> str:
     try:
         yt = YouTube(link)
         stream = yt.streams.get_audio_only()
@@ -25,13 +42,14 @@ def download_audio(link):
         return audio_path
     except Exception as e:
         print(f"Unable to download file. Exception {e}")
-        return None
+        return ""
+
 
-    
 demo = gr.Interface(
     fn=infer,
-    inputs= "text",
-    outputs= "text"
+    inputs=[gr.Textbox(label = "Youtube Link", placeholder="Copy link here"), gr.Checkbox(value=True, label="Add timestamps?")],
+    outputs=[gr.Textbox(label = "Transcription", placeholder="Should be here after almost same time as video length")],
+    examples=[ ["https://www.youtube.com/watch?v=KL2T0XRzWUI", False], ["https://www.youtube.com/watch?v=yGB_K_xlHdI", False], ["https://www.youtube.com/watch?v=dv9sgFHS2Do", True],]
 )
 
-demo.launch()
+demo.launch()