Update app.py

app.py

@@ -12,9 +12,16 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Function to generate text
 def generate_text(prompt, max_new_tokens=200):
-    inputs = tokenizer(prompt, return_tensors="pt")
+    # Tokenize the input prompt
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+    
+    # Create an attention mask
+    attention_mask = inputs.attention_mask
+
+    # Generate the text using the model
     outputs = model.generate(
         inputs.input_ids,
+        attention_mask=attention_mask,  # Pass the attention mask
         max_new_tokens=max_new_tokens,
         do_sample=True,
         top_k=50,  # Consider only the top 50 tokens for sampling
@@ -22,6 +29,8 @@ def generate_text(prompt, max_new_tokens=200):
         temperature=0.8,  # Increase diversity slightly
         pad_token_id=tokenizer.eos_token_id  # Ensure correct behavior for padding
     )
+    
+    # Decode the output text
     text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
     return text
 
@@ -67,7 +76,57 @@ def update_dropdowns(stage):
     situation_choices = situations[stage]
     return gr.Dropdown.update(choices=product_choices), gr.Dropdown.update(choices=situation_choices)
 
-with gr.Blocks() as demo:
+# Custom CSS for design
+custom_css = """
+#app {
+    background-color: white;
+    color: #333;
+}
+
+.gradio-container {
+    font-family: 'Helvetica Neue', sans-serif;
+}
+
+.gr-button {
+    background-color: #ff4240;
+    color: white;
+    border: none;
+    padding: 10px 20px;
+    font-weight: bold;
+    transition: background-color 0.3s ease;
+}
+
+.gr-button:hover {
+    background-color: #e03a3a;
+}
+
+.gr-button-secondary {
+    background-color: #3dd2ce;
+    color: white;
+    border: none;
+    padding: 10px 20px;
+    font-weight: bold;
+    transition: background-color 0.3s ease;
+}
+
+.gr-button-secondary:hover {
+    background-color: #37bebc;
+}
+
+.gr-textbox, .gr-dropdown {
+    border: 2px solid #3dd2ce;
+    border-radius: 5px;
+    padding: 8px;
+}
+
+.gr-textbox:focus, .gr-dropdown:focus {
+    border-color: #ff4240;
+    box-shadow: 0 0 5px rgba(255, 66, 64, 0.5);
+}
+"""
+
+# Gradio Interface with CSS styling
+with gr.Blocks(css=custom_css) as demo:
     gr.Markdown("# Herramienta de Generación de Copies Publicitarios")
 
     with gr.Row():
@@ -83,7 +142,7 @@ with gr.Blocks() as demo:
     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")
+    generate_custom_button = gr.Button("Generar Copies Personalizados", elem_classes="gr-button-secondary")
     custom_output_text = gr.Textbox(label="Copies Personalizados")
 
     generate_custom_button.click(fn=generate_text, inputs=custom_prompt_input, outputs=custom_output_text)