Update app.py

app.py

@@ -9,7 +9,7 @@ from nltk.tokenize import sent_tokenize
 import gradio as gr
 import warnings
 import torch
-from transformers import pipeline, WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoProcessor, AutoModelForSpeechSeq2Seq, AutoTokenizer, AutoModelForSeq2SeqLM
 from pydub import AudioSegment
 import soundfile as sf
 import numpy as np
@@ -24,15 +24,13 @@ warnings.filterwarnings("ignore")
 HF_AUTH_TOKEN = os.getenv('HF_AUTH_TOKEN')
 
 
-tokenizer = WhisperTokenizer.from_pretrained("NbAiLabBeta/nb-whisper-medium")
-model = WhisperForConditionalGeneration.from_pretrained("NbAiLabBeta/nb-whisper-medium")
-processor = WhisperProcessor.from_pretrained("NbAiLabBeta/nb-whisper-medium")
-
+model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
+processor = AutoProcessor.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 
-asr = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, device=device, torch_dtype=torch.float32)
+asr = pipeline("automatic-speech-recognition", model=model, processor=processor, device=device, torch_dtype=torch.float32)
 
 def transcribe_audio(audio_file):
     with torch.no_grad():
@@ -82,9 +80,9 @@ def transcribe_audio(audio_file, batch_size=4):
         inputs = inputs.to(device)
         attention_mask = inputs.attention_mask.to(device) if 'attention_mask' in inputs else None
         with torch.no_grad():
-            output = transcription_model.generate(
+            output = model.generate(
                 inputs.input_features,
-                max_length=2048,  # Increase max_length for longer outputs
+                max_length=2048, 
                 num_beams=7,
                 task="transcribe",
                 attention_mask=attention_mask,
@@ -123,10 +121,7 @@ def summarize_text(text):
 # HTML syntax for imagery
 image_html = """
 <div style="text-align: center;">
-    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/Olas%20AudioSwitch%20Shop.png" alt="Banner" width="87%" height="auto">
-</div>
-<div style="text-align: center; margin-top: 20px;">
-    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/picture.jpg" alt="Additional Image" width="68%" height="auto">
+    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/picture.png" alt="Banner" width="87%" height="auto">
 </div>
 """
 
@@ -135,11 +130,11 @@ iface = gr.Blocks()
 
 with iface:
     gr.HTML(image_html)
-    gr.Markdown("# Switch Work Audio Transcription App\nUpload an audio file to get the transcription")
+    gr.Markdown("# Upload an audio file to get the transcription")
     audio_input = gr.Audio(type="filepath")
     batch_size_input = gr.Slider(minimum=1, maximum=16, step=1, label="Batch Size")
-    transcription_output = gr.Textbox()
-    summary_output = gr.Textbox()
+    transcription_output = gr.Textbox("Transcription | nb-whisper-large-semantic")
+    summary_output = gr.Textbox("Summary | TextRank, graph-based")
     transcribe_button = gr.Button("Transcribe and Summarize")
 
     def transcribe_and_summarize(audio_file, batch_size):