Update app_flash.py

app_flash.py

@@ -1,6 +1,7 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline as hf_pipeline
+from transformers import AutoTokenizer, pipeline
 from flashpack.integrations.transformers import FlashPackTransformersModelMixin
+from transformers import AutoModelForCausalLM
 
 # ============================================================
 # 1️⃣ FlashPack-enabled model class
@@ -9,32 +10,32 @@ class FlashPackGemmaModel(AutoModelForCausalLM, FlashPackTransformersModelMixin)
     pass
 
 # ============================================================
-# 2️⃣ Model & tokenizer loading
+# 2️⃣ Model & tokenizer settings
 # ============================================================
 MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
-FLASHPACK_REPO = "rahul7star/FlashPack"  # Upload target repo
+FLASHPACK_REPO = "rahul7star/FlashPack"
 
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# ============================================================
+# 3️⃣ Load or create FlashPack model
+# ============================================================
 try:
-    # Try loading directly from the FlashPack repo
     print("📂 Loading model from FlashPack repository...")
     model = FlashPackGemmaModel.from_pretrained_flashpack(FLASHPACK_REPO)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-except Exception as e:
-    print(f"⚠️ Could not load FlashPack model: {e}")
-    print("⚙️ Loading from HF Hub and saving FlashPack to the repository...")
-
+except FileNotFoundError:
+    print("⚠️ FlashPack model not found on Hub. Creating and uploading...")
     # Load from HF Hub
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     model = FlashPackGemmaModel.from_pretrained(MODEL_ID)
-
-    # Save directly to the Hugging Face repo
+    # Save as FlashPack directly to Hub
     model.save_pretrained_flashpack(FLASHPACK_REPO, push_to_hub=True)
-    print(f"✅ Model uploaded to Hugging Face Hub: {FLASHPACK_REPO}")
+    print(f"✅ Model uploaded as FlashPack to Hugging Face Hub: {FLASHPACK_REPO}")
 
 # ============================================================
-# 3️⃣ Text-generation pipeline
+# 4️⃣ Text-generation pipeline
 # ============================================================
-pipe = hf_pipeline(
+pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
@@ -42,18 +43,18 @@ pipe = hf_pipeline(
 )
 
 # ============================================================
-# 4️⃣ Prompt enhancement function
+# 5️⃣ Prompt enhancement function
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
-    # Build chat-template messages
+    # Build messages
     messages = [
         {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
         {"role": "user", "content": user_prompt},
     ]
 
-    # Apply tokenizer chat template
+    # Apply chat-template
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
     # Generate output
@@ -65,13 +66,13 @@ def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     )
     enhanced = outputs[0]["generated_text"].strip()
 
-    # Append to chat history
+    # Update chat history
     chat_history.append({"role": "user", "content": user_prompt})
     chat_history.append({"role": "assistant", "content": enhanced})
     return chat_history
 
 # ============================================================
-# 5️⃣ Gradio UI
+# 6️⃣ Gradio UI
 # ============================================================
 with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
@@ -95,7 +96,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 UI actions
+    # Bind 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)
@@ -110,7 +111,7 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
     )
 
 # ============================================================
-# 6️⃣ Launch
+# 7️⃣ Launch
 # ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)