app.py

import gradio as gr
import torch
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer, AwqConfig
import spaces

# Esperar a que una GPU esté disponible
def wait_for_gpu():
    while not torch.cuda.is_available():
        print("Esperando por una GPU...")
        time.sleep(10)
    print("GPU encontrada, continuando con la configuración.")

wait_for_gpu()

# Configuración del modelo y tokenizador usando AutoAWQ
@spaces.GPU(duration=120)  # Asegura la asignación de GPU durante la carga del modelo
def cargar_modelo():
    model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
    quantization_config = AwqConfig(
        bits=4,
        fuse_max_seq_len=512,
        do_fuse=True,
    )
    
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoAWQForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        device_map="auto",  # Distribuir automáticamente en las GPUs disponibles
        quantization_config=quantization_config
    )
    return model, tokenizer

# Cargar el modelo y el tokenizador
model, tokenizer = cargar_modelo()

# Instrucciones para la mejora de resoluciones judiciales
instructions = """
**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")

    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]

# Interfaz de usuario con Gradio
with gr.Blocks() as demo:
    chatbot = gr.ChatInterface(fn=mejorar_resolucion, title="Mejora de Resoluciones Judiciales")
    demo.launch()