Update app.py

app.py

@@ -3,28 +3,45 @@
 import gradio as gr
 import warnings
 import torch
-from transformers import WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor
+#from transformers import WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
 import soundfile as sf
 import ffmpeg
 import os
-from PIL import Image
-from huggingface_hub import InferenceClient
-from gradio_client import Client, file
-import spaces
+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")
 
-# Load tokenizer, model, and processor
-tokenizer = WhisperTokenizer.from_pretrained("NbAiLabBeta/nb-whisper-medium")
-model = WhisperForConditionalGeneration.from_pretrained("NbAiLabBeta/nb-whisper-medium")
-processor = WhisperProcessor.from_pretrained("NbAiLabBeta/nb-whisper-medium")
+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
+}
+
+
+processor = AutoProcessor.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
+model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
 
-# Set up device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 torch_dtype = torch.float32
 
-# Move model to device
 model.to(device)
 
 def convert_audio_format(audio_path):
@@ -32,12 +49,11 @@ def convert_audio_format(audio_path):
     ffmpeg.input(audio_path).output(output_path, format='wav', ar='16000').run(overwrite_output=True)
     return output_path
 
-# @spaces.GPU(duration=120, queue=False)
 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 * 28  # 28 seconds chunks
+    chunk_size = 16000 * 30  
     chunks = [audio_input[i:i + chunk_size] for i in range(0, len(audio_input), chunk_size)]
 
     transcription = ""
@@ -51,6 +67,7 @@ def transcribe_audio(audio_file, batch_size=4):
                 inputs.input_features,
                 max_length=2048,
                 num_beams=7,
+                task="transcribe",
                 attention_mask=attention_mask
             )
         transcription += " ".join(processor.batch_decode(output, skip_special_tokens=True)) + " "
@@ -61,8 +78,28 @@ def transcribe_audio(audio_file, batch_size=4):
 
     result = f"Transcription: {transcription.strip()}\n\nTime taken: {transcription_time:.2f} seconds\nNumber of words: {word_count}"
     
-    return result
-
+    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(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
 banner_html = """
@@ -72,7 +109,7 @@ banner_html = """
 """
 image_html = """
 <div style="text-align: center; margin-top: 20px;">
-    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/blob/main/500x_picture.png" alt="picture" width="50%" height="auto">
+    <img src="https://huggingface.co/spaces/camparchimedes/ola_s-audioshop/raw/main/500x_picture.png" alt="picture" width="50%" height="auto">
 </div>
 """
 
@@ -84,10 +121,18 @@ with iface:
     gr.Markdown("# 𝐍𝐯𝐢𝐝𝐢𝐚 𝐀𝟏𝟎𝟎 👋🏼👾🦾⚡ @{NbAiLab/whisper-norwegian-medium}\nUpload audio file: ☕")
     audio_input = gr.Audio(type="filepath")
     batch_size_input = gr.Slider(minimum=1, maximum=16, step=1, label="Batch Size")
-    transcription_output = gr.Textbox()
+    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")
 
-    transcribe_button.click(fn=transcribe_audio, inputs=[audio_input, batch_size_input], outputs=transcription_output)
+    def process_audio(audio_file, batch_size):
+        transcription, result = transcribe_audio(audio_file, batch_size)
+        pdf_path = save_to_pdf(transcription)
+        summary = summarize_text(transcription)
+        return result, pdf_path, summary
+
+    transcribe_button.click(fn=process_audio, inputs=[audio_input, batch_size_input], outputs=[transcription_output, pdf_output, summary_output])
 
 # Launch interface
 iface.launch(share=True, debug=True)