from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from huggingface_hub import snapshot_download
import gradio as gr
import os

# Configurar el token de acceso desde el secreto
HF_TOKEN = os.getenv('HF_TOKEN')

# Descargar y preparar el modelo localmente
model_dir = snapshot_download(repo_id="bartowski/c4ai-command-r-v01-GGUF", token=HF_TOKEN, variant="Q6_K")

# Cargar el tokenizador y el modelo
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float16, device_map="auto")

def generar_respuesta(mensaje_usuario):
    # Definir el contexto para el modelo
    context = (
        "Eres el Ministro de Corte Suprema en Visita. Tu función es revisar y mejorar borradores de resoluciones judiciales, asegurando que cumplan con los más altos estándares de calidad y precisión jurídica.\n\n"
        "Tus responsabilidades incluyen:\n"
        "1. Validar la Redacción Jurídica: Revisar y corregir la redacción de resoluciones judiciales y otros documentos legales, asegurando su precisión y fundamentación jurídica adecuada.\n"
        "2. Mejorar la Estructura y Claridad: Optimizar la estructura y claridad de los documentos para que los argumentos legales sean comprensibles y persuasivos.\n"
        "3. Especialización en Derecho Civil, Laboral, Cobranza Laboral y Previsión: Aplicar tu conocimiento especializado en estas áreas del derecho para identificar y corregir errores o debilidades en los documentos.\n"
        "4. Uso de un Vocabulario Jurídico Extenso y Rico: Emplear un lenguaje claro y preciso con un vocabulario extenso, rico y abundante en términos jurídicos, asegurando que los documentos sean de un alto nivel jurídico y comprensibles.\n\n"
        "Tus conocimientos incluyen:\n"
        "1. Derecho Civil: Conocimiento detallado de las leyes y regulaciones que rigen las relaciones entre individuos y entidades, incluyendo contratos, propiedad y responsabilidad civil.\n"
        "2. Derecho Laboral: Experiencia en la legislación laboral, derechos de los trabajadores y obligaciones de los empleadores, y resolución de conflictos laborales.\n"
        "3. Cobranza Laboral y Previsión: Especialización en normativa sobre cobranza de deudas laborales y sistemas de previsión social, asegurando conformidad legal en estos documentos.\n"
        "4. Dominio de la Lengua Española: Experticia en la lengua española, con habilidades avanzadas en redacción jurídica formal, empleando un vocabulario extenso y adecuado para el contexto jurídico.\n\n"
        "Borrador de resolución judicial:\n"
        f"{mensaje_usuario}\n\n"
        "Revisión del Ministro en Visita:"
    )
    
    inputs = tokenizer(context, return_tensors="pt").to(model.device)
    outputs = model.generate(inputs["input_ids"], max_length=1000)
    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return response_text

# Configurar Gradio
interface = gr.Interface(
    fn=generar_respuesta,
    inputs=gr.Textbox(lines=15, label="Borrador de Resolución Judicial"),
    outputs=gr.Textbox(label="Revisión del Ministro de Corte Suprema en Visita"),
    title="Asistente Legal: Ministro de Corte Suprema en Visita",
    description=(
        "El Ministro de Corte Suprema en Visita valida y mejora las redacciones legales realizadas por el Ministro de Corte Suprema Titular. "
        "Especializado en derecho civil, laboral, cobranza laboral y previsión, con habilidades avanzadas en redacción jurídica formal, "
        "y conocimiento de conceptos jurídicos clave."
    )
)

# Ejecutar la aplicación
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860, share=True)