Update app.py

app.py

@@ -3,57 +3,84 @@
 import gradio as gr
 import warnings
 import torch
-#from transformers import WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
+from transformers import pipeline, WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor
+
+warnings.filterwarnings("ignore")
+
+# Load tokenizer and model
+tokenizer = WhisperTokenizer.from_pretrained("NbAiLabBeta/nb-whisper-medium")
+model = WhisperForConditionalGeneration.from_pretrained("NbAiLabBeta/nb-whisper-medium")
+processor = WhisperProcessor.from_pretrained("NbAiLabBeta/nb-whisper-medium")
+
+# Set up the device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+torch_dtype = torch.float32
+
+# Initialize pipeline
+asr = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, device=device, torch_dtype=torch_dtype)
+
+def transcribe_audio(audio_file):
+    # Perform transcription
+    with torch.no_grad():
+        output = asr(audio_file, chunk_length_s=28, generate_kwargs={"num_beams": 5, "task": "transcribe", "language": "no"})
+    return output["text"]
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=transcribe_audio,
+    inputs=gr.Audio(source="upload", type="filepath"),
+    outputs="text",
+    title="Audio Transcription App",
+    description="Upload an audio file to get the transcription",
+    theme="default",
+    layout="vertical",
+    live=False
+)import gradio as gr
+import warnings
+import torch
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoTokenizer, AutoModelForSeq2SeqLM
+from pydub import AudioSegment
 import soundfile as sf
-import ffmpeg
+import numpy as np
 import os
+import nltk
 from fpdf import FPDF
 import time
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.cluster import KMeans
-import re
-
-import nltk
-from nltk.tokenize import word_tokenize
-from nltk.corpus import stopwords
-import pandas as pd
-
-warnings.filterwarnings("ignore")
 
 nltk.download('punkt')
-nltk.download('stopwords')
-
-#tokenizer = WhisperTokenizer.from_pretrained("NbAiLabBeta/nb-whisper-large")
-#model = WhisperForConditionalGeneration.from_pretrained("NbAiLabBeta/nb-whisper-large")
-#processor = WhisperProcessor.from_pretrained("NbAiLabBeta/nb-whisper-large")
-
-generation_config = {
-    "temperature": 0.8,
-    "top_p": 0.9,
-    "top_k": 0.5,
-    "max_output_tokens": 2048
-}
 
+warnings.filterwarnings("ignore")
 
+# Load processor and model for transcription
 processor = AutoProcessor.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
-model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
+transcription_model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
+
+# Load tokenizer and model for summarization
+summarization_tokenizer = AutoTokenizer.from_pretrained("NbAiLab/norbert-summarization")
+summarization_model = AutoModelForSeq2SeqLM.from_pretrained("NbAiLab/norbert-summarization")
 
+# Set up the device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 torch_dtype = torch.float32
 
-model.to(device)
+# Move the models to the device
+transcription_model.to(device)
+summarization_model.to(device)
 
-def convert_audio_format(audio_path):
-    output_path = "converted_audio.wav"
-    ffmpeg.input(audio_path).output(output_path, format='wav', ar='16000').run(overwrite_output=True)
-    return output_path
+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
 
 def transcribe_audio(audio_file, batch_size=4):
     start_time = time.time()
-    audio_path = convert_audio_format(audio_file)
-    audio_input, sample_rate = sf.read(audio_path)
-    chunk_size = 16000 * 30  
+    # Convert .m4a to .wav
+    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 = ""
@@ -63,12 +90,14 @@ 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 = model.generate(
+            output = transcription_model.generate(
                 inputs.input_features,
-                max_length=2048,
+                max_length=2048,  # Increase max_length for longer outputs
                 num_beams=7,
                 task="transcribe",
-                attention_mask=attention_mask
+                attention_mask=attention_mask,
+                # forced_decoder_ids=None,  # OBS! forced_decoder_ids must not be set. Just marked it out for, just in case..
+                language="no"
             )
         transcription += " ".join(processor.batch_decode(output, skip_special_tokens=True)) + " "
 
@@ -80,59 +109,68 @@ def transcribe_audio(audio_file, batch_size=4):
     
     return transcription.strip(), result
 
-def save_to_pdf(transcription):
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-    pdf.multi_cell(0, 10, transcription)
-    pdf_output_path = "transcription.pdf"
-    pdf.output(pdf_output_path)
-    return pdf_output_path
+    def summarize_text(text):
+    inputs = summarization_tokenizer([text], max_length=1024, return_tensors="pt", truncation=True)
+    inputs = inputs.to(device)
+    summary_ids = summarization_model.generate(inputs.input_ids, num_beams=4, max_length=150, early_stopping=True)
+    summary = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
 
-def summarize_text(transcription):
-    sentences = transcription.split(". ")
-    vectorizer = TfidfVectorizer(stop_words='norwegian')
-    X = vectorizer.fit_transform(sentences)
-    
-    kmeans = KMeans(n_clusters=1)
-    kmeans.fit(X)
-    avg = X.mean(axis=0)
-    summary = [sentences[i] for i in kmeans.predict(avg)]
-    
-    return ". ".join(summary) + "."
-
-# HTML
+    # HTML for banner and additional image
 banner_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>
-"""
-image_html = """
 <div style="text-align: center; margin-top: 20px;">
-    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/500x_picture.png" alt="picture" width="50%" height="auto">
+    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/picture.jpg" alt="Additional Image" width="50%" height="auto">
 </div>
 """
 
-# Gradio interface
+# Create Gradio interface
 iface = gr.Blocks()
 
 with iface:
     gr.HTML(banner_html)
-    gr.Markdown("# 𝐍𝐯𝐢𝐝𝐢𝐚 𝐀𝟏𝟎𝟎 👋🏼👾🦾⚡ @{NbAiLab/whisper-norwegian-medium}\nUpload audio file: ☕")
+    gr.Markdown("# Switch Work Audio Transcription App\nUpload 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(label="Transcription")
-    pdf_output = gr.File(label="Download Transcription as PDF")
-    summary_output = gr.Textbox(label="Summary")
-    transcribe_button = gr.Button("Transcribe")
-
-    def process_audio(audio_file, batch_size):
-        transcription, result = transcribe_audio(audio_file, batch_size)
-        pdf_path = save_to_pdf(transcription)
+    batch_size_input = gr.Slider(minimum=1, maximum=16, step=1, default=4, label="Batch Size")
+    transcription_output = gr.Textbox()
+    summary_output = gr.Textbox()
+    transcribe_button = gr.Button("Transcribe and Summarize")
+
+    def transcribe_and_summarize(audio_file, batch_size):
+        transcription = transcribe_audio(audio_file, batch_size)
         summary = summarize_text(transcription)
-        return result, pdf_path, summary
+        return transcription, summary
 
-    transcribe_button.click(fn=process_audio, inputs=[audio_input, batch_size_input], outputs=[transcription_output, pdf_output, summary_output])
+    transcribe_button.click(fn=transcribe_and_summarize, inputs=[audio_input, batch_size_input], outputs=[transcription_output, summary_output])
+    
 
-# Launch interface
+def save_to_pdf(transcription, summary):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+    
+    # Add transcription
+    pdf.multi_cell(0, 10, "Transcription:\n" + transcription)
+    
+    # Add a space between transcription and summary
+    pdf.ln(10)
+    
+    # Add summary
+    pdf.multi_cell(0, 10, "Summary:\n" + summary)
+    
+    pdf_output_path = "transcription_summary.pdf"
+    pdf.output(pdf_output_path)
+    return pdf_output_path
+
+
+    
+# Launch the interface
 iface.launch(share=True, debug=True)
+
+# Launch the interface
+iface.launch(share=True, debug=True)
+
+
+