Update app.py

app.py

@@ -16,14 +16,14 @@ login(token=hf_api_token, add_to_git_credential=True)
 
 
 # Define the local directory where the model will be saved
-local_model_dir = "./llama-2-7b-hf"
+local_model_dir = "./LlamaGuard-7b"
 
 # Create the directory if it doesn't exist
 os.makedirs(local_model_dir, exist_ok=True)
 
 # Download the model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir=local_model_dir)
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir=local_model_dir)
+model = AutoModelForCausalLM.from_pretrained("meta-llama/LlamaGuard-7b", cache_dir=local_model_dir)
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/LlamaGuard-7b", cache_dir=local_model_dir)
 
 # Initialize the Whisper processor and model
 whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
@@ -31,7 +31,7 @@ whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-
 
 # Initialize the summarization model and tokenizer
 # Load LLAMA 7B model with accelerate from local path
-model_name = "./llama-2-7b-hf"
+model_name = "./LlamaGuard-7b"
 with init_empty_weights():
     summarization_model = AutoModelForCausalLM.from_pretrained(model_name)