import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import spaces

# Instrucciones específicas para el modelo
instrucciones = """
**Instrucciones para GPT Personalizado "GPT Civil"**

... (detalles de las instrucciones) ...
"""

# 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)

# Especifica un pad_token_id si no está ya definido y es el mismo que eos_token_id
if tokenizer.pad_token_id is None or tokenizer.pad_token_id == tokenizer.eos_token_id:
    tokenizer.pad_token_id = tokenizer.eos_token_id + 1  # O cualquier otro valor que consideres adecuado

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{input_text}"
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
    attention_mask = inputs['attention_mask']

    outputs = model.generate(
        inputs.input_ids, 
        attention_mask=attention_mask, 
        max_new_tokens=500, 
        temperature=0.3, 
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id
    )
    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.")
    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()