Update app_low.py

app_low.py

@@ -1,47 +1,55 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # ============================================================
 # 1️⃣ Load model and tokenizer
 # ============================================================
 MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
 
-# Use CPU-friendly settings
+# Use GPU if available
 device = 0 if torch.cuda.is_available() else -1
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
-# Text-generation pipeline
 pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=device,       # 0 for GPU, -1 for CPU
+    device=device,  # 0 for GPU, -1 for CPU
 )
 
 # ============================================================
-# 2️⃣ Define the generation function
+# 2️⃣ Define the generation function (chat-template style)
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
-    full_prompt = f"Enhance and expand the following prompt with more details and context: {user_prompt}"
+    # Build messages using proper roles
+    messages = [
+        {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
+        {"role": "user", "content": user_prompt}
+    ]
 
+    # Use tokenizer chat template to build the input
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    # Generate output
     output = pipe(
-        full_prompt,
+        prompt,
         max_new_tokens=int(max_tokens),
         temperature=float(temperature),
         do_sample=True,
-    )[0]['generated_text'].strip()
+    )[0]["generated_text"].strip()
 
-    # Convert to Gradio messages format
+    # Append conversation to history
     chat_history.append({"role": "user", "content": user_prompt})
     chat_history.append({"role": "assistant", "content": output})
 
     return chat_history
 
+
 # ============================================================
 # 3️⃣ Gradio UI
 # ============================================================
@@ -49,7 +57,8 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
     gr.Markdown(
         """
         # ✨ Prompt Enhancer (Gemma 3 270M)
-        Enter a short prompt, and the model will expand it with extra details, context, and creativity.
+        Enter a short prompt, and the model will **expand it with details and creative context**  
+        using the Gemma chat-template interface.
         """
     )
 
@@ -66,7 +75,7 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
             send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
             clear_btn = gr.Button("🧹 Clear Chat")
 
-    # Bind functions
+    # Bind UI actions
     send_btn.click(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     user_prompt.submit(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
@@ -74,9 +83,9 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
     gr.Markdown(
         """
         ---
-        💡 Tips:
-        - Works best with short, descriptive prompts (e.g., "A cat sitting on a chair").
-        - Adjust temperature for creativity: higher = more diverse output.
+        💡 **Tips:**
+        - Works best with short, descriptive prompts (e.g., "a cat sitting on a chair")
+        - Increase *Temperature* for more creative output.
         """
     )