Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image
 import os
 
 # Set up the Hugging Face API key from environment variables
-hf_api_key = os.getenv("new_hf_token")
+hf_api_key = os.getenv("HF_API_KEY")
 if not hf_api_key:
     raise ValueError("Hugging Face API key not found! Please set the 'HF_API_KEY' environment variable.")
 headers = {"Authorization": f"Bearer {hf_api_key}"}
@@ -20,10 +20,14 @@ translation_model_name = "facebook/mbart-large-50-many-to-one-mmt"
 tokenizer = MBart50Tokenizer.from_pretrained(translation_model_name)
 translation_model = MBartForConditionalGeneration.from_pretrained(translation_model_name)
 
-# Load a text generation model from Hugging Face
+# Load a text generation model from Hugging Face using accelerate for memory optimization
 text_generation_model_name = "EleutherAI/gpt-neo-2.7B"  # You can switch to "EleutherAI/gpt-j-6B" if available
 text_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
-text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name, device_map="auto", torch_dtype=torch.float32)
+text_model = AutoModelForCausalLM.from_pretrained(
+    text_generation_model_name, 
+    device_map="auto",  # Automatically allocate model layers to devices (requires accelerate)
+    torch_dtype=torch.float32  # Specify dtype to optimize memory usage
+)
 
 # Create a pipeline for text generation
 text_generator = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer)