Update app.py

app.py

@@ -1,35 +1,22 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from huggingface_hub import login
-import os
 
-def init_huggingface():
-    """Initialize Hugging Face authentication either from secrets or user input"""
-    if 'HUGGING_FACE_TOKEN' not in st.session_state:
-        # First try to get from environment variable
-        token = os.getenv('HUGGINGFACE_TOKEN')
-        
-        # If not in environment, check streamlit secrets
-        if not token and 'huggingface_token' in st.secrets:
-            token = st.secrets['huggingface_token']
-            
-        # If still not found, ask user
-        if not token:
-            token = st.text_input('Enter your Hugging Face token:', type='password')
-            if not token:
-                st.warning('Please enter your Hugging Face token to proceed')
-                st.stop()
-        
-        st.session_state['HUGGING_FACE_TOKEN'] = token
-    
-    # Login to Hugging Face
-    login(st.session_state['HUGGING_FACE_TOKEN'])
-    return True
+# Verificar GPU al inicio
+def check_gpu():
+    if torch.cuda.is_available():
+        gpu_info = {
+            "GPU Disponible": True,
+            "Nombre GPU": torch.cuda.get_device_name(0),
+            "Memoria Total (GB)": round(torch.cuda.get_device_properties(0).total_memory/1e9, 2),
+            "CUDA Version": torch.version.cuda
+        }
+        return gpu_info
+    return {"GPU Disponible": False}
 
 class LlamaDemo:
     def __init__(self):
-        self.model_name = "meta-llama/Llama-2-70b-chat-hf"
+        self.model_name = "meta-llama/Llama-2-7b-chat-hf"
         self._model = None
         self._tokenizer = None
         
@@ -38,10 +25,9 @@ class LlamaDemo:
         if self._model is None:
             self._model = AutoModelForCausalLM.from_pretrained(
                 self.model_name,
-                torch_dtype=torch.float16,
+                torch_dtype=torch.float16,  # Usar float16 para optimizar memoria
                 device_map="auto",
-                trust_remote_code=True,
-                load_in_8bit=True  # Para optimizar memoria
+                load_in_8bit=True  # Cuantización 8-bit para optimizar memoria
             )
         return self._model
     
@@ -55,7 +41,6 @@ class LlamaDemo:
         return self._tokenizer
 
     def generate_response(self, prompt: str, max_new_tokens: int = 512) -> str:
-        # Format prompt for Llama 2 chat
         formatted_prompt = f"[INST] {prompt} [/INST]"
         
         inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
@@ -71,21 +56,26 @@ class LlamaDemo:
                 pad_token_id=self.tokenizer.eos_token_id
             )
         
+        # Liberar memoria GPU después de generar
+        torch.cuda.empty_cache()
+        
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response.split("[/INST]")[-1].strip()
 
 def main():
     st.set_page_config(
-        page_title="Llama 2 Demo",
+        page_title="Llama 2 Chat Demo",
         page_icon="🦙",
         layout="wide"
     )
     
     st.title("🦙 Llama 2 Chat Demo")
     
-    # Initialize Hugging Face authentication
-    if init_huggingface():
-        st.success("Successfully authenticated with Hugging Face!")
+    # Mostrar información de GPU
+    gpu_info = check_gpu()
+    with st.expander("💻 GPU Info", expanded=False):
+        for key, value in gpu_info.items():
+            st.write(f"{key}: {value}")
     
     # Initialize model
     if 'llama' not in st.session_state:
@@ -123,6 +113,14 @@ def main():
                         st.error(f"Error: {str(e)}")
     
     with st.sidebar:
+        st.markdown("""
+        ### Memory Management
+        To optimize GPU usage and costs:
+        - Model runs in 8-bit precision
+        - Memory is cleared after each generation
+        - Space sleeps after inactivity
+        """)
+        
         if st.button("Clear Chat History"):
             st.session_state.chat_history = []
             st.experimental_rerun()