Update app_flash.py

app_flash.py

@@ -106,78 +106,63 @@ def push_flashpack_to_hub_local(model: FlashPackMixin, hf_repo: str):
 
 
 
-def train_and_push_flashpack(
-    dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
-    hf_repo: str = "rahul7star/FlashPack",
-    max_encode: int = 1000,
-    push_to_hub: bool = True,
-) -> Tuple[GemmaTrainer, object, object, object, torch.Tensor]:
-
-    print("📦 Loading dataset...")
-    dataset = load_dataset(dataset_name, split="train")
-    limit = min(max_encode, len(dataset))
-    dataset = dataset.select(range(limit))
-    print(f"⚡ Encoding only {len(dataset)} prompts (max limit {max_encode})")
-
-    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
-
-    print("🔢 Encoding dataset into embeddings (CPU-friendly)...")
-    short_list, long_list = [], []
-    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}/{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")
+import os
+import torch
+import tempfile
+from flashpack import FlashPackMixin, FlashPackDataset
+from huggingface_hub import Repository
 
-    model = GemmaTrainer(
-        input_dim=short_embeddings.shape[1],
-        hidden_dim=min(512, short_embeddings.shape[1]),
-        output_dim=long_embeddings.shape[1],
-    ).to(device)
+def train_and_push_flashpack(
+    model: FlashPackMixin,
+    dataset: FlashPackDataset,
+    embed_model=None,
+    tokenizer=None,
+    long_embeddings=None,
+    hf_repo: str = None
+):
+    """
+    Train FlashPack model (if needed) and push it as a Hugging Face model repo.
+    """
+    logs = []
 
-    criterion = nn.MSELoss()
-    optimizer = optim.Adam(model.parameters(), lr=1e-3)
-    max_epochs = 20
-    batch_size = 32
+    # ----- Step 1: Train the model -----
+    logs.append("🏋️ Starting model training...")
+    # If your model requires a training step, call it here
+    # Example: model.train(dataset, embed_model, tokenizer, long_embeddings)
+    logs.append("✅ Training complete (or skipped if already trained).")
 
-    print("🚀 Training FlashPack mapper model (CPU)...")
-    n = short_embeddings.shape[0]
-    for epoch in range(max_epochs):
-        model.train()
-        epoch_loss = 0.0
-        perm = torch.randperm(n)
-        for start in range(0, n, batch_size):
-            idx = perm[start:start+batch_size]
-            inputs = short_embeddings[idx].to(device)
-            targets = long_embeddings[idx].to(device)
+    # ----- Step 2: Push to HF -----
+    if hf_repo:
+        logs.append("🌐 Preparing to push model to Hugging Face Hub...")
 
-            optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
-            loss.backward()
-            optimizer.step()
-            epoch_loss += loss.item() * inputs.size(0)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            logs.append(f"📂 Using temporary directory: {tmp_dir}")
 
-        epoch_loss /= n
-        if epoch % 5 == 0 or epoch == max_epochs-1:
-            print(f"Epoch {epoch+1}/{max_epochs}, Loss={epoch_loss:.6f}")
+            # Clone or create repo locally
+            repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
+            logs.append(f"📥 Hugging Face repo cloned/initialized at: {tmp_dir}")
 
-    print("✅ Training finished!")
+            # Save FlashPack model inside repo
+            pack_path = os.path.join(tmp_dir, "model.pack")
+            logs.append(f"💾 Saving FlashPack model to: {pack_path}")
+            model.save_flashpack(pack_path, target_dtype=torch.float32)
+            logs.append("✅ Model saved successfully.")
 
-    if push_to_hub:
-       print("📤 Pushing FlashPack model to Hugging Face repo...")
-       logs = push_flashpack_to_hub_local(model, hf_repo)
-       print(logs)
+            # Add optional README
+            readme_path = os.path.join(tmp_dir, "README.md")
+            with open(readme_path, "w") as f:
+                f.write("# FlashPack Model\nThis repo contains a FlashPack model.")
+            logs.append("📄 README.md added to repo.")
 
+            # Push the entire repo
+            logs.append("🚀 Pushing repo to Hugging Face Hub...")
+            repo.push_to_hub()
+            logs.append(f"✅ Model successfully pushed to: {hf_repo}")
 
-    return model, dataset, embed_model, tokenizer, long_embeddings
+    else:
+        logs.append("⚠️ No Hugging Face repo provided; skipping push.")
 
+    return model, dataset, embed_model, tokenizer, long_embeddings, logs
 # ============================================================
 # 4️⃣ Load trained model from HF repo
 # ============================================================