Update app.py

app.py

@@ -1,57 +1,65 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from huggingface_hub import login
+import os
 
-# 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}
-
-# Configurar autenticación
-def setup_auth():
-    if 'HUGGINGFACE_TOKEN' in st.secrets:
-        login(st.secrets['HUGGINGFACE_TOKEN'])
+def setup_llama3_auth():
+    """Configurar autenticación para Llama 3"""
+    if 'HUGGING_FACE_TOKEN_3' in st.secrets:
+        token = st.secrets['HUGGING_FACE_TOKEN_3']
+        login(token)
         return True
     else:
-        st.error("No se encontró el token de Hugging Face en los secrets")
+        st.error("No se encontró el token de Llama 3 en los secrets")
         st.stop()
         return False
 
-class LlamaDemo:
+class Llama3Demo:
     def __init__(self):
-        self.model_name = "meta-llama/Llama-2-7b-chat"
+        # Verificar autenticación antes de cargar el modelo
+        setup_llama3_auth()
+        
+        # Usando el modelo de 3B con instrucciones
+        self.model_name = "meta-llama/Llama-3.2-3B-Instruct"
         self._model = None
         self._tokenizer = None
         
     @property
     def model(self):
         if self._model is None:
-            self._model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                torch_dtype=torch.float16,  # Usar float16 para optimizar memoria
-                device_map="auto",
-                load_in_8bit=True  # Cuantización 8-bit para optimizar memoria
-            )
+            try:
+                self._model = AutoModelForCausalLM.from_pretrained(
+                    self.model_name,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                    load_in_8bit=True,  # Optimización de memoria
+                    use_auth_token=st.secrets['HUGGING_FACE_TOKEN_3']
+                )
+            except Exception as e:
+                st.error(f"Error cargando el modelo: {str(e)}")
+                st.error("Verifica tu acceso a Llama 3.2 en https://huggingface.co/meta-llama")
+                raise e
         return self._model
     
     @property
     def tokenizer(self):
         if self._tokenizer is None:
-            self._tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name,
-                trust_remote_code=True
-            )
+            try:
+                self._tokenizer = AutoTokenizer.from_pretrained(
+                    self.model_name,
+                    use_auth_token=st.secrets['HUGGING_FACE_TOKEN_3']
+                )
+            except Exception as e:
+                st.error(f"Error cargando el tokenizer: {str(e)}")
+                raise e
         return self._tokenizer
 
     def generate_response(self, prompt: str, max_new_tokens: int = 512) -> str:
-        formatted_prompt = f"[INST] {prompt} [/INST]"
+        # Formato específico para Llama 3.2
+        formatted_prompt = f"""<|system|>You are a helpful AI assistant.</s>
+<|user|>{prompt}</s>
+<|assistant|>"""
         
         inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
         
@@ -62,77 +70,81 @@ class LlamaDemo:
                 num_return_sequences=1,
                 temperature=0.7,
                 do_sample=True,
-                top_p=0.9,
-                pad_token_id=self.tokenizer.eos_token_id
+                top_p=0.9
             )
-        
-        # Liberar memoria GPU después de generar
-        torch.cuda.empty_cache()
+            
+            # Limpiar memoria GPU
+            torch.cuda.empty_cache()
         
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response.split("[/INST]")[-1].strip()
+        # Extraer solo la respuesta del asistente
+        return response.split("<|assistant|>")[-1].strip()
 
 def main():
-    st.set_page_config(
-        page_title="Llama 2 Chat Demo",
-        page_icon="🦙",
-        layout="wide"
-    )
+    st.set_page_config(page_title="Llama 3.2 Chat", page_icon="🦙")
     
-    st.title("🦙 Llama 2 Chat Demo")
+    st.title("🦙 Llama 3.2 Chat")
     
-    # 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}")
+    # Verificar configuración
+    with st.expander("🔧 Status", expanded=True):
+        try:
+            token_status = setup_llama3_auth()
+            st.write("Token Llama 3:", "✅" if token_status else "❌")
+            
+            if torch.cuda.is_available():
+                st.write("GPU:", torch.cuda.get_device_name(0))
+                st.write("Memoria GPU:", f"{torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
+            else:
+                st.warning("GPU no disponible")
+        except Exception as e:
+            st.error(f"Error en configuración: {str(e)}")
     
-    # Initialize model
+    # Inicializar el modelo
     if 'llama' not in st.session_state:
-        with st.spinner("Loading Llama 2... This might take a few minutes..."):
-            st.session_state.llama = LlamaDemo()
+        with st.spinner("Inicializando Llama 3.2... esto puede tomar unos minutos..."):
+            try:
+                st.session_state.llama = Llama3Demo()
+            except Exception as e:
+                st.error("Error inicializando el modelo")
+                st.stop()
     
-    if 'chat_history' not in st.session_state:
-        st.session_state.chat_history = []
-        
-    # Chat interface
-    with st.container():
-        for message in st.session_state.chat_history:
-            with st.chat_message(message["role"]):
-                st.write(message["content"])
+    # Gestión del historial de chat
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
     
-        if prompt := st.chat_input("What would you like to discuss?"):
-            st.session_state.chat_history.append({
-                "role": "user",
-                "content": prompt
-            })
-            
-            with st.chat_message("user"):
-                st.write(prompt)
-            
-            with st.chat_message("assistant"):
-                with st.spinner("Thinking..."):
-                    try:
-                        response = st.session_state.llama.generate_response(prompt)
-                        st.write(response)
-                        st.session_state.chat_history.append({
-                            "role": "assistant",
-                            "content": response
-                        })
-                    except Exception as e:
-                        st.error(f"Error: {str(e)}")
+    # Mostrar historial
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
     
+    # Interface de chat
+    if prompt := st.chat_input("Escribe tu mensaje aquí"):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        with st.chat_message("assistant"):
+            try:
+                response = st.session_state.llama.generate_response(prompt)
+                st.markdown(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+            except Exception as e:
+                st.error(f"Error generando respuesta: {str(e)}")
+    
+    # Sidebar con información y controles
     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
+        ### Acerca de
+        Este demo usa Llama 3.2-3B-Instruct, el nuevo modelo de Meta.
+        
+        ### Características
+        - Modelo de 3B parámetros
+        - Optimizado para diálogo
+        - Cuantización de 8-bits
         """)
         
-        if st.button("Clear Chat History"):
-            st.session_state.chat_history = []
+        if st.button("Limpiar Chat"):
+            st.session_state.messages = []
             st.experimental_rerun()
 
 if __name__ == "__main__":