Update app.py

app.py

@@ -1,5 +1,5 @@
 """
-Version: 4th_pruned_optimized_transcription_app.py
+Version: 5th_pruned_optimized_transcription_app.py (alias HF_modded_nb-whisper_T4)
 
 Description: webapp, transkribering (norsk), NbAiLab/nb-whisper-large, oppsummering, pdf-download.
 """
@@ -22,7 +22,7 @@ import warnings
 from pydub import AudioSegment
 import torch
 import torchaudio
-from transformers import pipeline
+from transformers import pipeline, WhisperTokenizer, WhisperForConditionalGeneration, WhisperProcessor
 from huggingface_hub import model_info
 import spacy
 import networkx as nx
@@ -46,35 +46,29 @@ def convert_to_wav(audio_file):
     audio.export(wav_file, format="wav")
     return wav_file
 
-
-# D3efine model
-MODEL_NAME = "NbAiLab/nb-whisper-large"
-lang = "no"
-
 # Initialize device for torch
 device = 0 if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float32
 
-# Define pipeline config
-pipe = pipeline(
-    task="automatic-speech-recognition",
-    model=MODEL_NAME,
-    chunk_length_s=30,
-    device=device,
-)
-
-
-
-# # Set eos_token_id and pad_token_id to different values
-#pipe.model.config.eos_token_id = 0
-#pipe.model.config.pad_token_id = 1
+# Load tokenizer and model
+tokenizer = WhisperTokenizer.from_pretrained("NbAiLab/nb-whisper-large")
+model = WhisperForConditionalGeneration.from_pretrained("NbAiLab/nb-whisper-large")
+processor = WhisperProcessor.from_pretrained("NbAiLab/nb-whisper-large")
 
-# Simple
-pipe.model.config.pad_token_id = pipe.tokenizer.get_decoder_prompt_ids(language=lang)
 
+# Model script does not support JIT compilation
+#model = model.to(device)
+#model = torch.jit.script(model)
 
-#assert pipe.model.config.eos_token_id != pipe.model.config.pad_token_id
-#"eos_token_id and pad_token_id must be different"
+# Generation kwargs
+generate_kwargs = {
+    "num_beams": 5,
+    "task": "transcribe",
+    "language": "no"
+}
 
+# Initialize pipeline
+asr = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, device=device, torch_dtype=torch_dtype)
 
 # Transcribe audio
 def transcribe_audio(audio_file):
@@ -83,7 +77,9 @@ def transcribe_audio(audio_file):
 
     start_time = time.time()
 
-    text = pipe(audio_file)["text"]
+    # ASR pipeline on audio
+    with torch.no_grad():
+        text = asr(audio_file, chunk_length_s=30, generate_kwargs=generate_kwargs)["text"]
 
     output_time = time.time() - start_time
 
@@ -112,7 +108,6 @@ def transcribe_audio(audio_file):
 
     return text, result
 
-
 # Clean and preprocess text for summarization
 def clean_text(text):
     text = re.sub(r'https?:\/\/.*[\r\n]*', '', text)
@@ -270,3 +265,6 @@ with iface:
             pdf_both.click(fn=lambda text, summary: save_to_pdf(text, summary), inputs=[text_output, summary_output], outputs=[pdf_output])
 
 iface.launch(share=True, debug=True)
+
+
+