app.py

import warnings
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# Ignore specific warnings
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers.tokenization_utils_base")

# Load model and tokenizer
model_name = "EleutherAI/gpt-neo-1.3B"  # Use a lightweight model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function to generate text
def generate_text(prompt, max_new_tokens=200):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_k=50,  # Consider only the top 50 tokens for sampling
        top_p=0.95,
        temperature=0.8,  # Increase diversity slightly
        pad_token_id=tokenizer.eos_token_id  # Ensure correct behavior for padding
    )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return text

# Function to generate situations and copy according to user selections
def generate_situations_and_copies(stage, product, situation):
    prompt = f"""
    Estás creando una campaña publicitaria para la etapa {stage} del funnel de marketing. El producto es: {product}.
    La situación cotidiana es: {situation}. Integra el concepto de marca "Tú eliges" en el copy.
    
    Genera 5 opciones de copy (titular y bajada) en español:
    1.
    2.
    3.
    4.
    5.
    """
    return generate_text(prompt)

# Dropdown options for Gradio
stages = ["Upper", "Middle", "Lower"]
products = {
    "Upper": ["Convierte tu teléfono en un POS", "Más de 700 plantillas web con Wix"],
    "Middle": ["Elige el modelo de dispositivo POS", "Entrega express en Lima"],
    "Lower": ["Servicios de cobro inmediato", "Cargas Ilimitadas de productos a Wix"]
}
situations = {
    "Upper": [
        "No elegiste dónde naciste, pero sí cómo gestionar tus ventas.",
        "No pudiste elegir el nombre que te pusieron, pero sí puedes elegir que tu teléfono sea tu POS."
    ],
    "Middle": [
        "No elegiste todas las decisiones difíciles, pero sí puedes elegir la herramienta perfecta para cada transacción.",
        "No puedes elegir todas las circunstancias, pero sí el modelo de POS adecuado para tu negocio."
    ],
    "Lower": [
        "No elegiste que el cliente pagara tarde, pero sí cómo cobrar rápido.",
        "No puedes elegir el ritmo de las ventas, pero sí cómo acelerar los cobros."
    ]
}

def update_dropdowns(stage):
    product_choices = products[stage]
    situation_choices = situations[stage]
    return gr.Dropdown.update(choices=product_choices), gr.Dropdown.update(choices=situation_choices)

with gr.Blocks() as demo:
    gr.Markdown("# Herramienta de Generación de Copies Publicitarios")

    with gr.Row():
        stage_input = gr.Dropdown(choices=stages, label="Etapa del Funnel", value="Upper")
        product_input = gr.Dropdown(choices=products["Upper"], label="Producto")
        situation_input = gr.Dropdown(choices=situations["Upper"], label="Situación")

    stage_input.change(fn=update_dropdowns, inputs=stage_input, outputs=[product_input, situation_input])

    generate_button = gr.Button("Generar Copies")
    output_text = gr.Textbox(label="Copies Generados")

    generate_button.click(fn=generate_situations_and_copies, inputs=[stage_input, product_input, situation_input], outputs=output_text)

    custom_prompt_input = gr.Textbox(label="Ingresa un prompt para personalizar el copy", value="")
    generate_custom_button = gr.Button("Generar Copies Personalizados")
    custom_output_text = gr.Textbox(label="Copies Personalizados")

    generate_custom_button.click(fn=generate_text, inputs=custom_prompt_input, outputs=custom_output_text)

demo.launch()