A

app.py

@@ -1,17 +1,9 @@
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
-from transformers import pipeline
 import spaces
 
-@spaces.GPU(duration=120)  # Solicita el uso de GPU por 120 segundos
-def inicializar_pipeline():
-    # Configuración del pipeline con el modelo
-    nombre_modelo = "mistralai/Mistral-Nemo-Instruct-2407"
-    chatbot = pipeline("text-generation", model=nombre_modelo, device=0)  # Use GPU
-    return chatbot
-
-# Inicializar el pipeline una vez
-chatbot = inicializar_pipeline()
-
 # Instrucciones específicas para el modelo
 instrucciones = """
 **Instrucciones para GPT Personalizado "GPT Civil"**
@@ -60,21 +52,35 @@ Como Ministro de la Corte Suprema de Chile, su objetivo es mejorar y finalizar l
    - Asegúrese de una presentación concisa y unificada de toda la información.
 """
 
-@spaces.GPU(duration=120)  # Solicita el uso de GPU por 120 segundos
-def mejorar_resolucion(texto_entrada):
+# Configuración del dispositivo
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Carga del modelo y el tokenizador
+model_name = "mistralai/Mistral-Nemo-Instruct-2407"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None
+)
+
+@spaces.GPU(duration=120)
+def mejorar_resolucion(input_text):
     # Construcción del prompt con instrucciones y entrada del usuario
-    prompt = f"{instrucciones}\n\n{texto_entrada}"
-    respuesta = chatbot(prompt, max_new_tokens=500, temperature=0.3, device=0)
-    return respuesta[0]['generated_text']
+    prompt = f"Deberas adoptar el rol y las instrucciones que se encuentran en {instrucciones}\n\n Y luego mejorar la redaccion de la resolución judicial o responder a la consulta realizada aca: {input_text}"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(inputs.input_ids, max_new_tokens=500, temperature=0.3, do_sample=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 # Definición de la interfaz de Gradio
 with gr.Blocks() as demo:
     gr.Markdown("# Mejora de Resoluciones Judiciales con GPT Civil")
     gr.Markdown("Utiliza el modelo Mistral-Nemo-Instruct-2407 para mejorar borradores de resoluciones judiciales.")
-    texto_entrada = gr.Textbox(label="Introduce tu resolución judicial")
-    texto_salida = gr.Textbox(label="Resolución mejorada")
-    boton_enviar = gr.Button("Mejorar")
-    boton_enviar.click(fn=mejorar_resolucion, inputs=texto_entrada, outputs=texto_salida)
+    input_text = gr.Textbox(label="Introduce tu resolución judicial")
+    output_text = gr.Textbox(label="Resolución mejorada")
+    submit_button = gr.Button("Mejorar")
+    submit_button.click(fn=mejorar_resolucion, inputs=input_text, outputs=output_text)
 
 # Lanzamiento de la aplicación
-demo.launch()
+demo.launch()