Update app.py

app.py

@@ -6,12 +6,24 @@ import gradio as gr
 chat_model_state = None
 chat_tokenizer_state = None
 
+# Inicialización de ZeroGPU (opcional)
+def initialize_zero_gpu():
+    """Inicializa ZeroGPU si es requerido por el entorno."""
+    try:
+        import spaces
+        spaces.GPU(lambda x: x)  # Realiza una inicialización dummy
+        print("ZeroGPU inicializado correctamente.")
+    except ImportError:
+        print("ZeroGPU no está disponible o no es necesario en este entorno.")
+
+# Llamamos a la inicialización de ZeroGPU al inicio
+initialize_zero_gpu()
+
 def load_chat_model():
     """Función para cargar el modelo de chat."""
     global chat_model_state, chat_tokenizer_state
     try:
         model_name = "Qwen/Qwen2.5-3B-Instruct"
-        print("Cargando el modelo de chat...")
         # Cargar el modelo en CPU o GPU según disponibilidad
         chat_model_state = AutoModelForCausalLM.from_pretrained(
             model_name,
@@ -22,13 +34,11 @@ def load_chat_model():
         print("Modelo cargado exitosamente.")
     except Exception as e:
         print(f"Error al cargar el modelo de chat: {e}")
-        chat_model_state = None
-        chat_tokenizer_state = None
 
-def generate_response(messages):
+def generate_response(messages, model, tokenizer):
     """Genera una respuesta usando el modelo de chat."""
     try:
-        if chat_model_state is None or chat_tokenizer_state is None:
+        if model is None or tokenizer is None:
             raise ValueError("El modelo de chat o el tokenizer no están cargados.")
         
         # Construir el prompt manualmente a partir del historial de mensajes
@@ -46,17 +56,17 @@ def generate_response(messages):
         prompt += "Assistant:"
 
         # Tokenizar el prompt
-        model_inputs = chat_tokenizer_state(prompt, return_tensors="pt").to(chat_model_state.device)
-        generated_ids = chat_model_state.generate(
+        model_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        generated_ids = model.generate(
             **model_inputs,
             max_new_tokens=512,
             temperature=0.7,
             top_p=0.95,
-            eos_token_id=chat_tokenizer_state.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
 
         # Decodificar la respuesta generada
-        generated_text = chat_tokenizer_state.decode(generated_ids[0], skip_special_tokens=True)
+        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
         
         # Extraer solo la respuesta del asistente
         response = generated_text[len(prompt):].strip()
@@ -74,12 +84,7 @@ with gr.Blocks() as app_chat:
     clear_btn_chat = gr.Button("Limpiar Conversación")
 
     conversation_state = gr.State(
-        value=[
-            {
-                "role": "system",
-                "content": "Eres un chatbot. Responde a las preguntas del usuario de manera concisa y clara.",
-            }
-        ]
+        value=[{"role": "system", "content": "Eres un chatbot. Responde a las preguntas del usuario de manera concisa y clara."}]
     )
 
     def process_input(text, history, conv_state):
@@ -91,7 +96,7 @@ with gr.Blocks() as app_chat:
         history.append((text, None))
 
         # Generar la respuesta del modelo de chat
-        response = generate_response(conv_state)
+        response = generate_response(conv_state, chat_model_state, chat_tokenizer_state)
 
         conv_state.append({"role": "assistant", "content": response})
         history[-1] = (text, response)