Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
 import os
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 # Retrieve the token from the environment variable
 hf_api_token = os.getenv("HF_API_TOKEN")
@@ -18,13 +19,19 @@ whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
 whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
 
 # Initialize the summarization model and tokenizer
-# Use a smaller version of the Llama model and load in FP16
-summarization_model = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/LlamaGuard-7b",
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True
+# Load LLAMA 7B model with accelerate
+model_name = "meta-llama/Llama-2-7b-hf"
+with init_empty_weights():
+    summarization_model = AutoModelForCausalLM.from_pretrained(model_name)
+
+summarization_model = load_checkpoint_and_dispatch(
+    summarization_model,
+    checkpoint_path=model_name,
+    device_map="auto",
+    dtype=torch.float16
 )
-summarization_tokenizer = AutoTokenizer.from_pretrained("meta-llama/LlamaGuard-7b")
+
+summarization_tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Function to transcribe audio
 def transcribe_audio(audio_file):
@@ -51,7 +58,7 @@ def process_audio(audio_file):
 # Gradio UI
 iface = gr.Interface(
     fn=process_audio,
-    inputs=gr.Audio(source="upload", type="file"),
+    inputs=gr.Audio(type="file"),
     outputs=[
         gr.Textbox(label="Transcription"),
         gr.Textbox(label="Summary")