Update app.py

app.py

@@ -1,31 +1,45 @@
 import gradio as gr
 import torch
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer, AutoModelForSpeechSeq2Seq
 
-# Whisper Model for Transcription
+# Whisper Model Optimization
 WHISPER_MODEL = "openai/whisper-large-v3"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+whisper_quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_quant_type="nf4"
+)
+
+whisper_model = AutoModelForSpeechSeq2Seq.from_pretrained(
+    WHISPER_MODEL,
+    device_map="auto",
+    quantization_config=whisper_quant_config
+)
+
+whisper_tokenizer = AutoTokenizer.from_pretrained(WHISPER_MODEL)
 transcriber = pipeline(
     task="automatic-speech-recognition",
-    model=WHISPER_MODEL,
+    model=whisper_model,
+    tokenizer=whisper_tokenizer,
     chunk_length_s=30,
-    device=DEVICE,
+    device=DEVICE
 )
 
-# LLaMA Model for Generating Meeting Minutes
-LLAMA = "meta-llama/Llama-2-7b-chat-hf"  # Change to your preferred model
-quant_config = BitsAndBytesConfig(
+# LLaMA Model Optimization
+LLAMA = "meta-llama/Llama-2-7b-chat-hf"
+llama_quant_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
     bnb_4bit_compute_dtype=torch.bfloat16,
     bnb_4bit_quant_type="nf4"
 )
 
-# Load Model & Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(LLAMA)
 tokenizer.pad_token = tokenizer.eos_token
-model = AutoModelForCausalLM.from_pretrained(LLAMA, device_map="auto", quantization_config=quant_config)
+model = AutoModelForCausalLM.from_pretrained(LLAMA, device_map="auto", quantization_config=llama_quant_config)
 
 # Function to Transcribe & Generate Minutes
 def process_audio(audio_file):