Update app.py

app.py

@@ -1,11 +1,20 @@
-# app.py
-
+import os
 import gradio as gr
 import warnings
 import torch
-from transformers import pipeline, WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor
+from transformers import WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor, AutoTokenizer, AutoModelForSeq2SeqLM
+from pydub import AudioSegment
+import soundfile as sf
+import numpy as np
+import nltk
+from fpdf import FPDF
+import time
 
 warnings.filterwarnings("ignore")
+nltk.download('punkt')
+
+# Load environment variable
+HF_AUTH_TOKEN = os.getenv('HF_AUTH_TOKEN')
 
 # Load tokenizer and model
 tokenizer = WhisperTokenizer.from_pretrained("NbAiLabBeta/nb-whisper-medium")
@@ -15,9 +24,10 @@ 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
+model.to(device)
 
 # Initialize pipeline
-asr = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, device=device, torch_dtype=torch_dtype)
+asr = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, device=device, torch_dtype=torch.float32)
 
 def transcribe_audio(audio_file):
     # Perform transcription
@@ -33,38 +43,16 @@ iface = gr.Interface(
     title="Audio Transcription App",
     description="Upload an audio file to get the transcription",
     theme="default",
+    layout="vertical",
     live=False
 )
 
-    
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoTokenizer, AutoModelForSeq2SeqLM
-from pydub import AudioSegment
-import soundfile as sf
-import numpy as np
-import os
-import nltk
-from fpdf import FPDF
-import time
-
-nltk.download('punkt')
-
-HF_AUTH_TOKEN = os.getenv('HF_AUTH_TOKEN')
-
-# transcription
-processor = AutoProcessor.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
-transcription_model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLabBeta/nb-whisper-large-semantic")
+# Load summarization models with authentication token
+summarization_tokenizer = AutoTokenizer.from_pretrained("t5-base")
+summarization_model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
 
-# summarization
-summarization_tokenizer = AutoTokenizer.from_pretrained("NbAiLab/norbert-summarization")
-summarization_model = AutoModelForSeq2SeqLM.from_pretrained("NbAiLab/norbert-summarization")
-
-# setup
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-torch_dtype = torch.float32
-
-# move 'em
-transcription_model.to(device)
-summarization_model.to(device)  # PS. model needs to be told to use graph-based summary method (Lexname?)
+# Move summarization model to device
+summarization_model.to(device)
 
 def convert_to_wav(audio_file):
     audio = AudioSegment.from_file(audio_file, format="m4a")
@@ -95,7 +83,6 @@ def transcribe_audio(audio_file, batch_size=4):
                 num_beams=7,
                 task="transcribe",
                 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)) + " "
@@ -109,7 +96,7 @@ def transcribe_audio(audio_file, batch_size=4):
     return transcription.strip(), result
 
 def summarize_text(text):
-    inputs = summarization_tokenizer([text], max_length=1024, return_tensors="pt", truncation=True)
+    inputs = summarization_tokenizer("summarize: " + text, return_tensors="pt", max_length=1024, 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)
@@ -144,7 +131,6 @@ with iface:
 
     transcribe_button.click(fn=transcribe_and_summarize, inputs=[audio_input, batch_size_input], outputs=[transcription_output, summary_output])
     
-
 def save_to_pdf(transcription, summary):
     pdf = FPDF()
     pdf.add_page()
@@ -163,7 +149,5 @@ def save_to_pdf(transcription, summary):
     pdf.output(pdf_output_path)
     return pdf_output_path
 
-
-    
 # run
 iface.launch(share=True, debug=True)