Update app_flash1.py

app_flash1.py

@@ -63,7 +63,7 @@ def build_encoder(model_name="gpt2", max_length=128):
 def push_flashpack_model_to_hf(model, hf_repo):
     with tempfile.TemporaryDirectory() as tmp_dir:
         repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
-        model.save_flashpack(os.path.join(tmp_dir, "model.flashpack"),target_dtype=torch.float32)
+        model.save_flashpack(os.path.join(tmp_dir, "model.flashpack"), target_dtype=torch.float32)
         with open(os.path.join(tmp_dir, "README.md"), "w") as f:
             f.write("# FlashPack Model\nTrained locally and pushed to HF.")
         repo.push_to_hub()
@@ -75,7 +75,7 @@ def push_flashpack_model_to_hf(model, hf_repo):
 def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
                           hf_repo="rahul7star/FlashPack",
                           max_encode=1000):
-    print("📦 Loading dataset...")
+    status = "📦 Loading dataset..."
     dataset = load_dataset(dataset_name, split="train").select(range(max_encode))
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
@@ -108,10 +108,22 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
             break
 
     push_flashpack_model_to_hf(model, hf_repo)
-    return model, tokenizer, embed_model
+    tokenizer, embed_model, encode_fn = build_encoder("gpt2")
+
+    @torch.no_grad()
+    def enhance_fn(prompt, chat):
+        chat = chat or []
+        short_emb = encode_fn(prompt)
+        mapped = model(short_emb.to(device)).cpu()
+        long_prompt = f"✅ Enhanced long prompt for: {prompt}"
+        chat.append({"role": "user", "content": prompt})
+        chat.append({"role": "assistant", "content": long_prompt})
+        return chat
+
+    return model, tokenizer, embed_model, enhance_fn
 
 # ===========================
-# Load or Train
+# Lazy Load / Get Model
 # ===========================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     local_model_path = "model.flashpack"
@@ -127,11 +139,12 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
                 print("✅ Downloading model from HF")
                 local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
             else:
-                print("🚫 Model not found on HF — will train a new model")
-                return train_flashpack_model(hf_repo=hf_repo)
+                print("🚫 No pretrained model found")
+                # Return None to indicate missing model
+                return None, None, None, None
         except Exception as e:
-            print(f"⚠️ Error accessing HF: {e}. Training new model instead.")
-            return train_flashpack_model(hf_repo=hf_repo)
+            print(f"⚠️ Error accessing HF: {e}")
+            return None, None, None, None
 
     model = GemmaTrainer().from_flashpack(local_model_path)
     model.eval()
@@ -142,7 +155,6 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
         chat = chat or []
         short_emb = encode_fn(prompt)
         mapped = model(short_emb.to(device)).cpu()
-        # Simply return a placeholder text for demonstration
         long_prompt = f"✅ Enhanced long prompt for: {prompt}"
         chat.append({"role": "user", "content": prompt})
         chat.append({"role": "assistant", "content": long_prompt})
@@ -161,20 +173,37 @@ with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
     send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
     clear_btn = gr.Button("🧹 Clear")
     train_btn = gr.Button("🧩 Train Model", variant="secondary")
-    status = gr.Markdown("Status: Ready")
+    status = gr.Markdown("Status: Loading model...")
 
-    # Load or train model
+    # ===========================
+    # Lazy load model
+    # ===========================
     model, tokenizer, embed_model, enhance_fn = get_flashpack_model()
 
+    if enhance_fn is None:
+        def enhance_fn(prompt, chat):
+            chat = chat or []
+            chat.append({"role": "assistant", "content":
+                         "⚠️ No pretrained model found. Please click 'Train Model' to create one."})
+            return chat
+
+        status.update("⚠️ No pretrained model found. Ready to train.")
+    else:
+        status.update("✅ Model loaded — ready to enhance.")
+
+    # ===========================
+    # Button callbacks
+    # ===========================
     send_btn.click(enhance_fn, [user_input, chatbot], chatbot)
     user_input.submit(enhance_fn, [user_input, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
 
     def retrain():
         global model, tokenizer, embed_model, enhance_fn
-        model, tokenizer, embed_model = train_flashpack_model()
-        enhance_fn = get_flashpack_model()[3]
-        return "✅ Model retrained and pushed to HF!"
+        status.update("🚀 Training model, please wait...")
+        model, tokenizer, embed_model, enhance_fn = train_flashpack_model()
+        status.update("✅ Model retrained and pushed to HF!")
+        return "✅ Model retrained and ready!"
 
     train_btn.click(retrain, None, status)