Update app.py

app.py

@@ -1,7 +1,7 @@
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelWithLMHead, TranslationPipeline
 import gradio as gr
-from pytube import YouTube
-import os
+
+
 
 pipe = pipeline(model="torileatherman/train_first_try")  # change to "your-username/the-name-you-picked"
 
@@ -9,28 +9,32 @@ def transcribe(audio):
     text = pipe(audio)["text"]
     return text
 
-def transcribe_url(url):
-    youtube = YouTube(str(url))
-    audio = youtube.streams.filter(only_audio=True).first().download('yt_video')
-    text = pipe(audio)["text"]
-    return text
+translation_pipeline = TranslationPipeline( model=AutoModelWithLMHead.from_pretrained("SEBIS/legal_t5_small_trans_sv_en"),
+                                            tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path = "SEBIS/legal_t5_small_trans_sv_en", 
+                                            do_lower_case=False, 
+                                            skip_special_tokens=True), 
+                                            device=0)
+
+def translate(text):
+    translation = translation_pipeline([text], max_length=512)
+    return translation
+
+demo = gr.Blocks()
+
+with demo:
+    
+    title="Whisper Small Swedish",
+    description="Realtime demo for Swedish speech recognition using a fine-tuned Whisper small model."
+    
+    inputs_audio = gr.Audio(source="microphone", type="filepath"), 
+ 
+    text = gr.Textbox()
+    translation = gr.Label()
+
+    b1 = gr.Button("Record audio")
+    b2 = gr.Button("Translate text")
 
-url_demo = gr.Interface(
-    fn=transcribe_url, 
-    inputs="text", 
-    outputs="text",
-    title="Whisper Swedish",
-    description="Swedish speech and audio recognition using a fine-tuned Whisper small model",
-)
-
-voice_demo = gr.Interface(
-    fn=transcribe, 
-    inputs=gr.Audio(source="microphone", type="filepath"), 
-    outputs="text",
-    title="Whisper Swedish",
-    description="Swedish speech and audio recognition using a fine-tuned Whisper small model",
-)
-
-demo = gr.TabbedInterface([url_demo, voice_demo], ["YouTube Video to Text", "Audio to Text"])
+    b1.click(transcribe, inputs=inputs_audio, outputs=text)
+    b2.click(translate, inputs=text, outputs=translation)
 
 demo.launch()