a

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 import torch
-from awq import AutoAWQForCausalLM
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
 class Modelo:
     def __init__(self):
-        self.model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
+        self.model_id = "mistralai/Mistral-Nemo-Instruct-2407"
         self.tokenizer = None
         self.model = None
         self.instrucciones = """
@@ -14,20 +13,27 @@ class Modelo:
 
 Como Ministro de la Corte Suprema de Chile, su objetivo es mejorar y finalizar los borradores de resoluciones judiciales para el sistema judicial civil del país...
 """
-        self.cargar_modelo_y_tokenizador()  # Llama a la función para inicializar el modelo y el tokenizador
-
+    
+    @spaces.GPU(duration=120)
     def cargar_modelo_y_tokenizador(self):
         if not torch.cuda.is_available():
             raise Exception("No hay GPU disponible")
+        
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
-        self.model = AutoAWQForCausalLM.from_pretrained(
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
             self.model_id,
-            torch_dtype=torch.float16,
-            device_map={"": "cuda:0"},  # Forzar la carga en GPU
+            torch_dtype=torch.bfloat16,  # Utiliza BF16 para ahorrar memoria
+            device_map="auto",  # Asigna automáticamente en función de los recursos disponibles
             low_cpu_mem_usage=True
         )
+        self.model.to("cuda")  # Mover explícitamente el modelo a la GPU si es posible
 
+    @spaces.GPU(duration=120)
     def mejorar_resolucion(self, message, history=None):
+        if self.model is None or self.tokenizer is None:
+            self.cargar_modelo_y_tokenizador()
+        
         prompt = [
             {"role": "system", "content": self.instrucciones},
             {"role": "user", "content": f"Mejora esta resolución judicial: {message}"}
@@ -47,4 +53,4 @@ modelo = Modelo()
 
 with gr.Blocks() as demo:
     chatbot = gr.ChatInterface(fn=modelo.mejorar_resolucion, title="Mejora de Resoluciones Judiciales")
-    demo.launch()
+    demo.launch()