a

app.py

@@ -1,9 +1,21 @@
 import gradio as gr
 from transformers import pipeline
-import spaces  # Importar el módulo spaces para usar el decorador @spaces.GPU
+import torch
+import spaces
+import time
 
-# Configuración del pipeline con el modelo
-pipe = pipeline("text-generation", model="mistralai/Mistral-Nemo-Instruct-2407")
+# Función para verificar y esperar una GPU
+def wait_for_gpu():
+    while not torch.cuda.is_available():
+        print("Esperando por una GPU...")
+        time.sleep(10)  # Espera 10 segundos antes de volver a verificar
+    print("GPU encontrada, continuando con la configuración.")
+
+# Espera hasta que una GPU esté disponible
+wait_for_gpu()
+
+# Configuración del pipeline con el modelo Meta-Llama usando GPU
+pipe = pipeline("text-generation", model="meta-llama/Meta-Llama-3.1-405B-Instruct-FP8", device=0)
 
 # Instrucciones para mejorar resoluciones judiciales
 instructions = """
@@ -52,13 +64,13 @@ Como Ministro de la Corte Suprema de Chile, su objetivo es mejorar y finalizar l
    - Proporcione una respuesta completa, incluyendo correcciones, sugerencias y la resolución final.
    - Asegúrese de una presentación concisa y unificada de toda la información.
 """
-@spaces.GPU(duration=120)  # Decorar la función con @spaces.GPU y especificar duración
-def mejorar_resolucion(message, history):
-    prompt = f"{instructions}\n\n{message}"
-    response = pipe(prompt, max_new_tokens=2000, return_full_text=False, clean_up_tokenization_spaces=True)
+@spaces.GPU(duration=120)  # Ajusta la duración según el tiempo estimado de la tarea
+def mejorar_resolucion(message, history=None):
+    prompt = f"Analiza profundamente y aplica estas {instructions}\n\n y luego mejora esta resolución judicial {message}"
+    response = pipe(prompt, max_new_tokens=1000, return_full_text=False, clean_up_tokenization_spaces=True)
     return response[0]["generated_text"]
 
 # Configuración de la interfaz de Gradio
 with gr.Blocks() as demo:
     chatbot = gr.ChatInterface(fn=mejorar_resolucion, title="Mejora de Resoluciones Judiciales")
-    demo.launch(share=True)
+    demo.launch()