Update app_flash.py

app_flash.py

@@ -1,52 +1,57 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM
 from flashpack.integrations.transformers import FlashPackTransformersModelMixin
+import os
 
 # ============================================================
 # 1️⃣ FlashPack-enabled model class
 # ============================================================
 class FlashPackGemmaModel(AutoModelForCausalLM, FlashPackTransformersModelMixin):
-    """AutoModelForCausalLM extended with FlashPackMixin for fast save/load"""
+    """AutoModelForCausalLM extended with FlashPackMixin for local save/load"""
     pass
 
-MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
 
 # ============================================================
-# 2️⃣ Load model and tokenizer with FlashPack
+# 2️⃣ Model and tokenizer setup
 # ============================================================
-try:
-    print("📂 Trying to load model from FlashPack directory...")
-    model = FlashPackGemmaModel.from_pretrained_flashpack("model_flashpack")
+MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
+FLASHPACK_DIR = "model_flashpack"
+
+if os.path.exists(FLASHPACK_DIR):
+    print("📂 Loading model from local FlashPack directory...")
+    model = FlashPackGemmaModel.from_pretrained_flashpack(FLASHPACK_DIR)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-except Exception as e:
-    print("⚙️ FlashPack model not found, loading from Hugging Face Hub...")
+else:
+    print("⚙️ Loading model from Hugging Face Hub...")
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    # Load Hugging Face model and wrap into FlashPack class
     model = FlashPackGemmaModel.from_pretrained(MODEL_ID)
-    # Save for future faster loads
-    model.save_pretrained_flashpack("model_flashpack")
-    print("✅ Model saved as FlashPack for next startup!")
+    # Save locally as FlashPack for next run
+    model.save_pretrained_flashpack(FLASHPACK_DIR, push_to_hub=False)
+    print("✅ Model saved locally as FlashPack!")
+
 
 # ============================================================
-# 3️⃣ Create text-generation pipeline
+# 3️⃣ Text-generation pipeline
 # ============================================================
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 
 
 # ============================================================
-# 4️⃣ Define prompt enhancement logic
+# 4️⃣ Prompt enhancement function
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
+    # Build chat-template messages
     messages = [
         {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
         {"role": "user", "content": user_prompt},
     ]
 
-    # Use chat-template
+    # Apply tokenizer chat-template
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
+    # Generate enhanced prompt
     outputs = pipe(
         prompt,
         max_new_tokens=int(max_tokens),
@@ -56,7 +61,7 @@ def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
 
     enhanced = outputs[0]["generated_text"].strip()
 
-    # Append to chat
+    # Append to chat history
     chat_history.append({"role": "user", "content": user_prompt})
     chat_history.append({"role": "assistant", "content": enhanced})
 
@@ -64,7 +69,7 @@ def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
 
 
 # ============================================================
-# 5️⃣ Gradio Interface
+# 5️⃣ Gradio UI
 # ============================================================
 with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft()) as demo:
     gr.Markdown(