Update app_flash.py

app_flash.py

@@ -1,21 +1,21 @@
 import gradio as gr
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer
 from flashpack.integrations.transformers import FlashPackTransformersModelMixin
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, pipeline as hf_pipeline
 
 # ============================================================
-# 1️⃣ FlashPack-enabled model class
+# 1️⃣ Define FlashPack-enabled model class
 # ============================================================
 class FlashPackGemmaModel(AutoModelForCausalLM, FlashPackTransformersModelMixin):
+    """Gemma 3 model wrapped with FlashPackTransformersModelMixin"""
     pass
 
 # ============================================================
-# 2️⃣ Model & tokenizer settings
+# 2️⃣ Load tokenizer
 # ============================================================
 MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
 FLASHPACK_REPO = "rahul7star/FlashPack"
 
-# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 # ============================================================
@@ -25,17 +25,15 @@ try:
     print("📂 Loading model from FlashPack repository...")
     model = FlashPackGemmaModel.from_pretrained_flashpack(FLASHPACK_REPO)
 except FileNotFoundError:
-    print("⚠️ FlashPack model not found on Hub. Creating and uploading...")
-    # Load from HF Hub
+    print("⚠️ FlashPack model not found. Loading from HF Hub and uploading FlashPack...")
     model = FlashPackGemmaModel.from_pretrained(MODEL_ID)
-    # Save as FlashPack directly to Hub
     model.save_pretrained_flashpack(FLASHPACK_REPO, push_to_hub=True)
-    print(f"✅ Model uploaded as FlashPack to Hugging Face Hub: {FLASHPACK_REPO}")
+    print(f"✅ FlashPack model uploaded to Hugging Face Hub: {FLASHPACK_REPO}")
 
 # ============================================================
-# 4️⃣ Text-generation pipeline
+# 4️⃣ Build text-generation pipeline
 # ============================================================
-pipe = pipeline(
+pipe = hf_pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
@@ -43,30 +41,26 @@ pipe = pipeline(
 )
 
 # ============================================================
-# 5️⃣ Prompt enhancement function
+# 5️⃣ Define prompt enhancement function
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
-    # Build messages
     messages = [
         {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
         {"role": "user", "content": user_prompt},
     ]
 
-    # Apply chat-template
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
-    # Generate output
     outputs = pipe(
         prompt,
         max_new_tokens=int(max_tokens),
         temperature=float(temperature),
-        do_sample=True,
+        do_sample=True
     )
     enhanced = outputs[0]["generated_text"].strip()
 
-    # Update chat history
     chat_history.append({"role": "user", "content": user_prompt})
     chat_history.append({"role": "assistant", "content": enhanced})
     return chat_history
@@ -96,7 +90,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 actions
+    # 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)