Update app.py

app.py

@@ -1,39 +1,35 @@
 
 
-import whisper
 import gradio as gr
+import whisper
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-
-# Initialize the device map for ZeRO
-from accelerate.utils import set_module_tensor_to_device
 import torch
 
-device_map = "auto"  # Automatically allocate layers across available GPUs/CPUs
+# Check if GPU is available and set up device map
+device_map = "auto"  # Automatically balance layers across available devices
 print(f"Using ZeRO-powered device map: {device_map}")
 
-# Load the model using ZeRO
-model_name = "openai/whisper-tiny"
+# Load the Whisper model using Accelerate with ZeRO
+model_name = "tiny"  # Change to "base", "small", etc., as needed
 
-# Load the Whisper model into ZeRO's memory-efficient mode
+print(f"Loading the Whisper model: {model_name} with ZeRO optimization...")
 with init_empty_weights():
-    whisper_model = whisper.load_model(model_name)
+    whisper_model = whisper.load_model(model_name)  # Load model structure without weights
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Load model with Accelerate/ZeRO
+# Dispatch the model across devices using ZeRO
 whisper_model = load_checkpoint_and_dispatch(
     whisper_model,
     device_map=device_map,
-    dtype=torch.float16  # Optional: Use mixed precision for further optimization
+    dtype=torch.float16  # Use mixed precision for efficiency
 )
 
+print("Model successfully loaded with ZeRO optimization!")
+
 # Define the transcription function
 def transcribe(audio):
     # Perform transcription using the Whisper model
     result = whisper_model.transcribe(audio)
-    return result['text']
+    return result["text"]
 
 # Create the Gradio interface
 demo = gr.Interface(
@@ -41,7 +37,7 @@ demo = gr.Interface(
     inputs=gr.Audio(source="microphone", type="filepath", label="Speak into the microphone"),  # Input audio
     outputs=gr.Textbox(label="Transcription"),  # Output transcription
     title="Whisper Speech-to-Text with ZeRO",  # Title of the interface
-    description="Record audio using your microphone and get a transcription using the Whisper model optimized by ZeRO."
+    description="Record audio using your microphone and get a transcription using the Whisper model optimized with ZeRO."
 )
 
 # Launch the Gradio interface