Update app_flash.py

app_flash.py

@@ -1,36 +1,38 @@
 import gradio as gr
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline as hf_pipeline
 from flashpack.integrations.transformers import FlashPackTransformersModelMixin
-from transformers import AutoModelForCausalLM, pipeline as hf_pipeline
 
 # ============================================================
-# 1️⃣ Define FlashPack wrapper for Gemma
+# 1️⃣ FlashPack-enabled model class
 # ============================================================
 class FlashPackGemmaModel(AutoModelForCausalLM, FlashPackTransformersModelMixin):
     pass
 
 # ============================================================
-# 2️⃣ Load model & tokenizer via FlashPack pipeline approach
+# 2️⃣ Model & tokenizer loading
 # ============================================================
 MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
-FLASHPACK_REPO = "rahul7star/FlashPack"
+FLASHPACK_REPO = "rahul7star/FlashPack"  # Upload target repo
 
-# Load model from FlashPack repository (Hub or local path)
 try:
+    # Try loading directly from the FlashPack repo
     print("📂 Loading model from FlashPack repository...")
-    # Load model directly via FlashPack
     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("⚙️ Falling back to HF model and saving FlashPack...")
+    print("⚙️ Loading from HF Hub and saving FlashPack to the repository...")
+
+    # Load from HF Hub
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     model = FlashPackGemmaModel.from_pretrained(MODEL_ID)
-    model.save_pretrained_flashpack("model_flashpack", push_to_hub=False)
-    print("✅ Saved FlashPack locally for next run")
+
+    # Save directly to the Hugging Face repo
+    model.save_pretrained_flashpack(FLASHPACK_REPO, push_to_hub=True)
+    print(f"✅ Model uploaded to Hugging Face Hub: {FLASHPACK_REPO}")
 
 # ============================================================
-# 3️⃣ Create text-generation pipeline
+# 3️⃣ Text-generation pipeline
 # ============================================================
 pipe = hf_pipeline(
     "text-generation",
@@ -40,12 +42,12 @@ pipe = hf_pipeline(
 )
 
 # ============================================================
-# 4️⃣ Define prompt enhancement function
+# 4️⃣ Prompt enhancement function
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
-    # Build messages with chat template
+    # Build chat-template messages
     messages = [
         {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
         {"role": "user", "content": user_prompt},