Update app_flash1.py

app_flash1.py

@@ -89,17 +89,16 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
 
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
-    def encode_dataset(ds):
-        s_list, l_list = [], []
-        for i, item in enumerate(ds):
-            s_list.append(encode_fn(item["short_prompt"]))
-            l_list.append(encode_fn(item["long_prompt"]))
-            if (i + 1) % 50 == 0:
-                log_fn(f"  → Encoded {i + 1}/{len(ds)}")
-                gc.collect()
-        return torch.vstack(s_list), torch.vstack(l_list)
-
-    short_emb, long_emb = encode_dataset(dataset)
+    # Only encode short+long embeddings
+    s_list, l_list = [], []
+    for i, item in enumerate(dataset):
+        s_list.append(encode_fn(item["short_prompt"]))
+        l_list.append(encode_fn(item["long_prompt"]))
+        if (i + 1) % 50 == 0:
+            log_fn(f"  → Encoded {i + 1}/{len(dataset)}")
+            gc.collect()
+    short_emb, long_emb = torch.vstack(s_list), torch.vstack(l_list)
+
     model = GemmaTrainer()
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
     loss_fn = nn.CosineSimilarity(dim=1)
@@ -125,13 +124,11 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
         chat = chat or []
         short_emb = encode_fn(prompt)
         mapped = model(short_emb.to(device)).cpu()
-        long_prompt = f"✅ Enhanced long prompt for: {prompt}"
+        long_prompt = f"🌟 Enhanced prompt: {prompt} (creatively expanded)"
         chat.append({"role": "user", "content": prompt})
         chat.append({"role": "assistant", "content": long_prompt})
         return chat
 
-    
-
     return model, tokenizer, embed_model, enhance_fn, logs
 
 # ===========================
@@ -140,11 +137,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:
@@ -157,24 +152,16 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
             print(f"⚠️ Error accessing HF: {e}")
             return None, None, None, None
 
-    # Load the model
     model = GemmaTrainer().from_flashpack(local_model_path)
     model.eval()
-
-    # Load encoder
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
-    # Enhancement function (without dataset)
     @torch.no_grad()
     def enhance_fn(prompt, chat):
         chat = chat or []
         short_emb = encode_fn(prompt).to(device)
         mapped = model(short_emb).cpu()
-
-        # Convert the model output tensor to a string representation for demonstration
-        # In practice, you could use a small language head on top of mapped embeddings
-        long_prompt = f"✅ Enhanced long prompt generated for: {prompt}"
-
+        long_prompt = f"🌟 Enhanced prompt: {prompt} (creatively expanded)"
         chat.append({"role": "user", "content": prompt})
         chat.append({"role": "assistant", "content": long_prompt})
         return chat
@@ -201,8 +188,8 @@ with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
     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."})
+            chat.append({"role": "assistant",
+                         "content": "⚠️ No pretrained model found. Please click 'Train Model' to create one."})
             return chat
         logs.append("⚠️ No pretrained model found. Ready to train.")
     else: