Update app.py

app.py

@@ -1,21 +1,43 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
+import torch
 import os
 
 # Authenticate with Hugging Face
-HF_API_TOKEN = os.getenv("ttt") 
+HF_API_TOKEN = os.getenv("ttt")  # Ensure 'ttt' is set in the environment
+if HF_API_TOKEN is None:
+    raise ValueError("Hugging Face API token ('ttt') not set in environment variables.")
 login(HF_API_TOKEN)
 
-# Load your model
+# Load your model with memory optimization
 model_name = "Ouiam123/Llama-2-7b-chat-finetune-tourism"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_API_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HF_API_TOKEN)
+
+# Check for GPU availability
+use_gpu = torch.cuda.is_available()
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_API_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",  # Automatically map layers to devices (CPU/GPU)
+        offload_folder="./offload",  # Offload parts of the model to disk if needed
+        load_in_8bit=use_gpu,  # Enable 8-bit precision for GPU
+        use_auth_token=HF_API_TOKEN,
+    )
+except Exception as e:
+    raise RuntimeError(f"Failed to load the model: {e}")
 
 # Define the response generation function
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
+    # Preprocess input to optimize memory
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+
+    # Move inputs to GPU if available
+    if use_gpu:
+        inputs = {key: value.to("cuda") for key, value in inputs.items()}
+
+    outputs = model.generate(**inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Create a Gradio interface
@@ -23,8 +45,10 @@ interface = gr.Interface(
     fn=generate_response,
     inputs="text",
     outputs="text",
-    title="Moroccan Tourism Chatbot"
+    title="Optimized Moroccan Tourism Chatbot",
+    description="Ask any questions about tourism in Morocco!",
 )
 
 # Launch the app
-interface.launch()
+if __name__ == "__main__":
+    interface.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))