Update app_flash.py

app_flash.py

@@ -70,7 +70,7 @@ def train_flashpack_model(
     dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
     model_name: str = "gpt2",
     max_length: int = 32,
-    subset_limit: int | None = 500,  # None = use full dataset
+    max_encode: int = 2000,  # maximum number of prompts to encode
     push_to_hub: bool = False,
     hf_repo: str = "rahul7star/FlashPack",
 ) -> tuple:
@@ -79,14 +79,12 @@ def train_flashpack_model(
     print("📦 Loading dataset...")
     dataset = load_dataset(dataset_name, split="train")
 
-    # Handle subset for quick CPU training
-    if subset_limit is None:
-        subset_limit = len(dataset)
+    # Limit dataset to max_encode prompts
+    limit = min(max_encode, len(dataset))
+    dataset = dataset.select(range(limit))
+    print(f"⚡ Encoding only {len(dataset)} prompts (max limit {max_encode})")
 
-    dataset = dataset.select(range(min(subset_limit, len(dataset))))
-    print(f"⚡ Using subset: {len(dataset)} examples for quick training")
-
-    # 2️⃣ Setup tokenizer and encoder
+    # 2️⃣ Setup tokenizer & encoder
     tokenizer, embed_model, encode_fn = build_encoder(model_name=model_name, max_length=max_length)
 
     # 3️⃣ Encode dataset (CPU-friendly)
@@ -95,28 +93,35 @@ def train_flashpack_model(
     for i, item in enumerate(dataset):
         short_list.append(encode_fn(item["short_prompt"]))
         long_list.append(encode_fn(item["long_prompt"]))
-        if (i + 1) % 50 == 0 or (i + 1) == len(dataset):
-            print(f"  → Encoded {i+1}/{len(dataset)} prompts")
+
+        # Exit early if we hit max_encode
+        if (i + 1) >= max_encode:
+            print(f"⚡ Reached max encode limit: {max_encode} prompts, stopping early.")
+            break
+
+        # Progress logging
+        if (i + 1) % 50 == 0:
+            print(f"  → Encoded {i+1}/{limit} prompts")
             gc.collect()
 
     short_embeddings = torch.vstack(short_list)
     long_embeddings = torch.vstack(long_list)
+    print(f"✅ Finished encoding {short_embeddings.shape[0]} prompts")
 
-    # 4️⃣ Initialize model
+    # 4️⃣ Initialize and train model (same as before)
     model = GemmaTrainer(
         input_dim=short_embeddings.shape[1],
-        hidden_dim=min(512, short_embeddings.shape[1]),  # smaller hidden dim for speed
+        hidden_dim=min(512, short_embeddings.shape[1]),
         output_dim=long_embeddings.shape[1],
     ).to(device)
 
-    # 5️⃣ Training loop
     criterion = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
-    max_epochs = 50      # fewer epochs
+    max_epochs = 50
     tolerance = 1e-4
-    batch_size = 32       # smaller batch size for CPU
+    batch_size = 32
 
-    print("🚀 Training FlashPack mapper model (fast, CPU)...")
+    print("🚀 Training FlashPack mapper model (CPU)...")
     n = short_embeddings.shape[0]
     for epoch in range(max_epochs):
         model.train()