d

app.py

@@ -4,45 +4,51 @@ from awq import AutoAWQForCausalLM
 from transformers import AutoTokenizer
 import spaces
 
-@spaces.GPU(duration=120)
-def cargar_modelo():
-    model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
-    
-    # Omite quantization_config si ya está incluido en el modelo
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoAWQForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True,
-        device_map=None  # Evita especificar manualmente el device_map
-    )
-    return model, tokenizer
-
-model, tokenizer = cargar_modelo()
-
-instructions = """
+class Modelo:
+    def __init__(self):
+        self.model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
+        self.tokenizer = None
+        self.model = None
+
+    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_id,
+            torch_dtype=torch.float16,
+            device_map={"": "cuda:0"},  # Forzar la carga en GPU
+            low_cpu_mem_usage=True
+        )
+
+    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}"}
+        ]
+        inputs = self.tokenizer.apply_chat_template(
+            prompt,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            return_dict=True,
+        ).to("cuda")  # Asegura que los tensores se procesen en GPU
+
+        outputs = self.model.generate(**inputs, do_sample=True, max_new_tokens=500)
+        return self.tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)[0]
+
+    def __init__(self):
+        self.instrucciones = """
 **Instrucciones para GPT Personalizado "GPT Civil"**
 
 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...
 """
 
-@spaces.GPU(duration=120)
-def mejorar_resolucion(message, history=None):
-    prompt = [
-        {"role": "system", "content": instructions},
-        {"role": "user", "content": f"Mejora esta resolución judicial: {message}"}
-    ]
-    inputs = tokenizer.apply_chat_template(
-        prompt,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_tensors="pt",
-        return_dict=True,
-    ).to("cuda")  # Asegura que los tensores se procesen en GPU
-
-    outputs = model.generate(**inputs, do_sample=True, max_new_tokens=500)
-    return tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)[0]
+modelo = Modelo()
 
 with gr.Blocks() as demo:
-    chatbot = gr.ChatInterface(fn=mejorar_resolucion, title="Mejora de Resoluciones Judiciales")
-    demo.launch()
+    chatbot = gr.ChatInterface(fn=modelo.mejorar_resolucion, title="Mejora de Resoluciones Judiciales")
+    demo.launch()