Update app.py

app.py

@@ -5,89 +5,57 @@ import nltk
 nltk.download('punkt')
 from nltk.tokenize import sent_tokenize
 
-
 import gradio as gr
 import warnings
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoProcessor, AutoModelForSpeechSeq2Seq # pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 from pydub import AudioSegment
-import soundfile as sf
-import numpy as np
 from fpdf import FPDF
 from PIL import Image
 import time
 import os
-# import spaces
 
 warnings.filterwarnings("ignore")
 
-# HF_AUTH_TOKEN = os.getenv('HF_AUTH_TOKEN')
-
-processor = AutoProcessor.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
-model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
-
-
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 torch_dtype = torch.float32
-model.to(device)
 
+# Initialize the ASR pipeline
+pipe = pipeline("automatic-speech-recognition", model="NbAiLabBeta/nb-whisper-large-semantic", device=device, torch_dtype=torch.float32)
+
+# Function to convert m4a files to wav
 def convert_to_wav(audio_file):
     audio = AudioSegment.from_file(audio_file, format="m4a")
     wav_file = "temp.wav"
     audio.export(wav_file, format="wav")
     return wav_file
 
-
-# @spaces.GPU(queue=True)
-
-def transcribe_audio(audio_file, batch_size=4):
-    start_time = time.time()
-
+# Transcription function using the ASR pipeline
+def transcribe_audio(audio_file):
     if audio_file.endswith(".m4a"):
         audio_file = convert_to_wav(audio_file)
 
-    audio_input, sample_rate = sf.read(audio_file)
-    chunk_size = 16000 * 30
-    chunks = [audio_input[i:i + chunk_size] for i in range(0, len(audio_input), chunk_size)]
-
-    transcription = ""
-
-    for i in range(0, len(chunks), batch_size):
-        batch_chunks = chunks[i:i + batch_size]
-        inputs = processor(batch_chunks, sampling_rate=16000, return_tensors="pt", padding=True)
-        inputs = inputs.to(device)
-
-        attention_mask = inputs.attention_mask.to(device) if 'attention_mask' in inputs else None
-
-        with torch.no_grad():
-            output = model.generate(
-                inputs.input_features,
-                max_length=2048,
-                num_beams=8,
-                attention_mask=attention_mask,
-                pad_token_id=processor.tokenizer.pad_token_id,
-                eos_token_id=processor.tokenizer.eos_token_id
-            )
+    start_time = time.time()
 
-        transcription += " ".join(processor.batch_decode(output, skip_special_tokens=True)) + " "
+    with torch.no_grad():
+        output = pipe(audio_file, chunk_length_s=30, generate_kwargs={"num_beams": 8, "task": "transcribe", "language": "no"})
 
+    transcription = output["text"]
     end_time = time.time()
-    transcription_time = end_time - start_time
+
+    output_time = end_time - start_time
     word_count = len(transcription.split())
 
-    result = f"Transcription: {transcription.strip()}\n\nTime taken: {transcription_time:.2f} seconds\nNumber of words: {word_count}"
+    result = f"Transcription: {transcription.strip()}\n\nTime taken: {output_time:.2f} seconds\nNumber of words: {word_count}"
 
     return transcription.strip(), result
 
-
-# summarization model
+# Summarization model setup
 summarization_tokenizer = AutoTokenizer.from_pretrained("t5-base")
 summarization_model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
-
-# t5-base to device
 summarization_model.to(device)
 
-# Graph-based summarization|TextRank
+# Graph-based summarization (TextRank)
 def summarize_text(text):
     sentences = sent_tokenize(text)
     if len(sentences) == 0:
@@ -102,48 +70,29 @@ def summarize_text(text):
 
     ranked_sentences = sorted(((scores[i], s) for i, s in enumerate(sentences)), reverse=True)
 
-    # Select top N sentences (e.g., 3 sentences for the summary)
     top_n = 3
     summary = " ".join([s for _, s in ranked_sentences[:top_n]])
     return summary
 
-# HTML syntax for imagery
-image_html = """
-<div style="text-align: center;">
-    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/picture.png" alt="Banner" width="85%" height="auto">
-</div>
-"""
-
+# Save transcription and summary to PDF
 def save_to_pdf(transcription, summary):
     pdf = FPDF()
     pdf.add_page()
     pdf.set_font("Arial", size=12)
 
-    # include transcription
-    pdf.multi_cell(0, 10, "Transcription:\n" + transcription)
+    if transcription:
+        pdf.multi_cell(0, 10, "Transcription:\n" + transcription)
 
-    # paragraph space
     pdf.ln(10)
 
-    # include summary
-    pdf.multi_cell(0, 10, "Summary:\n" + summary)
+    if summary:
+        pdf.multi_cell(0, 10, "Summary:\n" + summary)
 
     pdf_output_path = "transcription_summary.pdf"
     pdf.output(pdf_output_path)
     return pdf_output_path
 
-# Gradio UI
-iface = gr.Interface(
-    fn=transcribe_audio,
-    inputs=gr.Audio(type="filepath"),
-    outputs="text",
-    title="Audio Transcription App",
-    description="Upload an audio file to get the transcription",
-    theme="default",
-    live=False
-)
-
-# Gradio UI
+# Gradio Interface setup
 iface = gr.Blocks()
 
 with iface:
@@ -151,26 +100,24 @@ with iface:
     gr.Markdown("# Vi har nå muligheten til å oversette lydfiler til norsk skrift.")
 
     with gr.Tabs():
-
-        # First Tab: Transcription
+        # Transcription Tab
         with gr.TabItem("Transcription"):
             audio_input = gr.Audio(type="filepath")
-            batch_size_input = gr.Slider(minimum=7, maximum=16, step=1, label="Batch Size")
             transcription_output = gr.Textbox(label="Transcription | nb-whisper-large-semantic")
             result_output = gr.Textbox(label="Time taken and Number of words")
             transcribe_button = gr.Button("Transcribe")
 
-            def transcribe(audio_file, batch_size):
-                transcription, result = transcribe_audio(audio_file, batch_size)
+            def transcribe(audio_file):
+                transcription, result = transcribe_audio(audio_file)
                 return transcription, result
 
             transcribe_button.click(
                 fn=transcribe,
-                inputs=[audio_input, batch_size_input],
+                inputs=[audio_input],
                 outputs=[transcription_output, result_output]
             )
 
-        # Second Tab: Summary
+        # Summary Tab
         with gr.TabItem("Summary"):
             summary_output = gr.Textbox(label="Summary | TextRank, graph-based")
             summarize_button = gr.Button("Summarize")
@@ -183,11 +130,11 @@ with iface:
 
             summarize_button.click(
                 fn=summarize,
-                inputs=[transcription_output],  # Use the transcription from the first tab
+                inputs=[transcription_output],
                 outputs=summary_output
             )
 
-        # Third Tab: PDF Download Options
+        # PDF Download Tab
         with gr.TabItem("Download PDF"):
             pdf_transcription_only = gr.Button("Download PDF with Transcription Only")
             pdf_summary_only = gr.Button("Download PDF with Summary Only")
@@ -224,8 +171,5 @@ with iface:
                 outputs=[pdf_output_both]
             )
 
-
-# run
+# Run the Gradio interface
 iface.launch(share=True, debug=True)
-
-