Create app_flash.py

app_flash.py

@@ -0,0 +1,107 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from flashpack.integrations.transformers import FlashPackTransformersModelMixin
+
+
+# ============================================================
+# 1️⃣ Extend model with FlashPackMixin
+# ============================================================
+class FlashPackedCausalLM(AutoModelForCausalLM, FlashPackTransformersModelMixin):
+    pass
+
+
+# ============================================================
+# 2️⃣ Load or prepare model
+# ============================================================
+MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
+
+try:
+    print("📂 Trying to load FlashPack model...")
+    model = FlashPackedCausalLM.from_pretrained_flashpack("model_flashpack")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+except Exception as e:
+    print("⚙️ FlashPack not found, loading from Hugging Face...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+    model.save_pretrained_flashpack("model_flashpack")
+    print("✅ Model saved as FlashPack for next startup!")
+
+# Create text-generation pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype="auto", device_map="auto")
+
+
+# ============================================================
+# 3️⃣ Define inference logic
+# ============================================================
+def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
+    chat_history = chat_history or []
+
+    messages = [
+        {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
+        {"role": "user", "content": user_prompt},
+    ]
+
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    outputs = pipe(
+        prompt,
+        max_new_tokens=int(max_tokens),
+        temperature=float(temperature),
+        do_sample=True,
+    )
+
+    enhanced = outputs[0]["generated_text"].strip()
+
+    # Add to chat
+    chat_history.append({"role": "user", "content": user_prompt})
+    chat_history.append({"role": "assistant", "content": enhanced})
+
+    return chat_history
+
+
+# ============================================================
+# 4️⃣ Gradio Interface
+# ============================================================
+with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # ✨ Prompt Enhancer (Gemma 3 270M)
+        Enter a short prompt, and the model will **expand it with details and creative context**  
+        using the Gemma chat-template interface.
+        """
+    )
+
+    with gr.Row():
+        chatbot = gr.Chatbot(height=400, label="Enhanced Prompts", type="messages")
+        with gr.Column(scale=1):
+            user_prompt = gr.Textbox(
+                placeholder="Enter a short prompt...",
+                label="Your Prompt",
+                lines=3,
+            )
+            temperature = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature")
+            max_tokens = gr.Slider(32, 256, value=128, step=16, label="Max Tokens")
+            send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
+            clear_btn = gr.Button("🧹 Clear Chat")
+
+    # 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)
+
+    gr.Markdown(
+        """
+        ---
+        💡 **Tips:**
+        - Works best with short, descriptive prompts (e.g., "a cat sitting on a chair")
+        - Increase *Temperature* for more creative output.
+        """
+    )
+
+
+# ============================================================
+# 5️⃣ Launch App
+# ============================================================
+if __name__ == "__main__":
+    demo.launch(show_error=True)