import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Configuración
base_model_id = "unsloth/gemma-2-2b-bnb-4bit"
adapter_id = "cochinoraptor/bloggon-ads-lora"

# Cargar el modelo base y el adaptador LoRA
print("Cargando modelo base...")
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    device_map="auto",
    torch_dtype=torch.float16
)

print("Cargando tokenizador...")
tokenizer = AutoTokenizer.from_pretrained(base_model_id)

print("Aplicando adaptador LoRA...")
model = PeftModel.from_pretrained(base_model, adapter_id)
model.eval()

def generate_blog(topic, temperature=0.7, max_tokens=400):
    prompt = f"Escribe un blog sobre {topic}:\n\nBlog:"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            do_sample=True,
            top_p=0.95,
            repetition_penalty=1.1
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Limpiar la respuesta para mostrar solo el blog
    if "Blog:" in response:
        response = response.split("Blog:")[-1].strip()
    return response

# Interfaz Gradio
iface = gr.Interface(
    fn=generate_blog,
    inputs=[
        gr.Textbox(label="Tema del blog", placeholder="Ejemplo: inteligencia artificial, viajes, marketing..."),
        gr.Slider(label="Creatividad (temperature)", minimum=0.1, maximum=1.5, value=0.7, step=0.1),
        gr.Slider(label="Longitud (tokens)", minimum=100, maximum=800, value=400, step=50)
    ],
    outputs=gr.Textbox(label="Blog generado", lines=20),
    title="Bloggon Pet - Generador de Blogs con IA",
    description="Ingresa un tema y la IA generará un blog completo para ti."
)

iface.launch()