Update app_flash1.py

app_flash1.py

@@ -60,14 +60,16 @@ def build_encoder(model_name="gpt2", max_length=128):
 # ===========================
 # Push model to HF
 # ===========================
-def push_flashpack_model_to_hf(model, hf_repo):
+def push_flashpack_model_to_hf(model, hf_repo, log_fn):
     with tempfile.TemporaryDirectory() as tmp_dir:
+        log_fn(f"📦 Preparing repository {hf_repo}...")
         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)
         with open(os.path.join(tmp_dir, "README.md"), "w") as f:
             f.write("# FlashPack Model\nTrained locally and pushed to HF.")
+        log_fn("⏳ Pushing model to Hugging Face...")
         repo.push_to_hub()
-        print(f"✅ Model pushed to {hf_repo}")
+        log_fn(f"✅ Model pushed to {hf_repo}")
 
 # ===========================
 # Training
@@ -75,8 +77,16 @@ 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):
-    status = "📦 Loading dataset..."
+    logs = []
+
+    def log_fn(msg):
+        logs.append(msg)
+        print(msg)
+    
+    log_fn("📦 Loading dataset...")
     dataset = load_dataset(dataset_name, split="train").select(range(max_encode))
+    log_fn(f"✅ Loaded {len(dataset)} samples")
+
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
     def encode_dataset(ds):
@@ -85,7 +95,7 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
             s_list.append(encode_fn(item["short_prompt"]))
             l_list.append(encode_fn(item["long_prompt"]))
             if (i + 1) % 50 == 0:
-                print(f"  → Encoded {i + 1}/{len(ds)}")
+                log_fn(f"  → Encoded {i + 1}/{len(ds)}")
                 gc.collect()
         return torch.vstack(s_list), torch.vstack(l_list)
 
@@ -94,7 +104,7 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
     loss_fn = nn.CosineSimilarity(dim=1)
 
-    print("🚀 Training model...")
+    log_fn("🚀 Training model...")
     for epoch in range(20):
         model.train()
         optimizer.zero_grad()
@@ -102,12 +112,12 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
         loss = 1 - loss_fn(preds, long_emb).mean()
         loss.backward()
         optimizer.step()
-        print(f"Epoch {epoch+1}/20 | Loss: {loss.item():.5f}")
+        log_fn(f"Epoch {epoch+1}/20 | Loss: {loss.item():.5f}")
         if loss.item() < 0.01:
-            print("🎯 Early stopping.")
+            log_fn("🎯 Early stopping.")
             break
 
-    push_flashpack_model_to_hf(model, hf_repo)
+    push_flashpack_model_to_hf(model, hf_repo, log_fn)
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
     @torch.no_grad()
@@ -120,7 +130,12 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
         chat.append({"role": "assistant", "content": long_prompt})
         return chat
 
-    return model, tokenizer, embed_model, enhance_fn
+    # Test model on sample prompt
+    test_prompt = "Hello world"
+    enhance_fn(test_prompt, [])
+    log_fn(f"✅ Model test complete: '{test_prompt}' -> Enhanced prompt available")
+
+    return model, tokenizer, embed_model, enhance_fn, logs
 
 # ===========================
 # Lazy Load / Get Model
@@ -128,11 +143,9 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     local_model_path = "model.flashpack"
 
-    # 1. Try local
     if os.path.exists(local_model_path):
         print("✅ Loading local model")
     else:
-        # 2. Try HF
         try:
             files = list_repo_files(hf_repo)
             if "model.flashpack" in files:
@@ -140,7 +153,6 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
                 local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
             else:
                 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}")
@@ -173,12 +185,11 @@ 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: Loading model...")
+    log_output = gr.Textbox(label="Logs", lines=15)
 
-    # ===========================
     # Lazy load model
-    # ===========================
     model, tokenizer, embed_model, enhance_fn = get_flashpack_model()
+    logs = []
 
     if enhance_fn is None:
         def enhance_fn(prompt, chat):
@@ -186,26 +197,23 @@ with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
             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.")
+        logs.append("⚠️ No pretrained model found. Ready to train.")
     else:
-        status.update("✅ Model loaded — ready to enhance.")
+        logs.append("✅ 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
-        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!"
+        global model, tokenizer, embed_model, enhance_fn, logs
+        logs = ["🚀 Training model, please wait..."]
+        model, tokenizer, embed_model, enhance_fn, train_logs = train_flashpack_model()
+        logs.extend(train_logs)
+        return gr.Textbox.update(value="\n".join(logs))
 
-    train_btn.click(retrain, None, status)
+    train_btn.click(retrain, None, log_output)
 
 if __name__ == "__main__":
     demo.launch(show_error=True)