Update app_flash.py

app_flash.py

@@ -60,6 +60,53 @@ def build_encoder(model_name="gpt2", max_length: int = 32):
 # ============================================================
 # 3️⃣ Train and push FlashPack model
 # ============================================================
+import os
+import tempfile
+from huggingface_hub import hf_hub_download, HfApi
+
+# ------------------------------------------------------------
+# Utility to push FlashPack model to HF using upload_file
+# ------------------------------------------------------------
+def push_flashpack_to_hub_local(model, hf_repo: str, size_limit_mb=100):
+    """
+    Saves FlashPack model locally, then uploads files to Hugging Face repo individually.
+    Avoids push_to_hub() call which may not be compatible.
+    """
+    logs = []
+
+    # 1️⃣ Save model locally to temp directory
+    temp_dir = tempfile.mkdtemp()
+    model.save_flashpack(temp_dir)
+    logs.append(f"📦 Model saved locally to {temp_dir}")
+
+    # 2️⃣ List saved files
+    files = [f for f in os.listdir(temp_dir) if os.path.isfile(os.path.join(temp_dir, f))]
+    logs.append(f"📄 Found {len(files)} files to upload.")
+
+    api = HfApi()
+    for i, file in enumerate(files, start=1):
+        file_path = os.path.join(temp_dir, file)
+        size_mb = os.path.getsize(file_path) / (1024*1024)
+        if size_mb > size_limit_mb:
+            logs.append(f"⏭️ Skipping {file}, size {size_mb:.2f}MB > limit {size_limit_mb}MB")
+            continue
+
+        logs.append(f"⬆️ Uploading {file} ({i}/{len(files)}) ...")
+        try:
+            api.upload_file(
+                path_or_fileobj=file_path,
+                path_in_repo=file,
+                repo_id=hf_repo,
+                repo_type="model",
+                commit_message=f"Add {file} from FlashPack training"
+            )
+            logs.append(f"✅ Uploaded {file}")
+        except Exception as e:
+            logs.append(f"⚠️ Failed to upload {file}: {e}")
+
+    logs.append(f"🎉 Model push complete to {hf_repo}")
+    return "\n".join(logs)
+
 def train_and_push_flashpack(
     dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
     hf_repo: str = "rahul7star/FlashPack",
@@ -125,10 +172,10 @@ def train_and_push_flashpack(
     print("✅ Training finished!")
 
     if push_to_hub:
-        print(f"📤 Pushing model to Hugging Face repo: {hf_repo} ...")
-        model.save_flashpack(hf_repo, target_dtype=torch.float32)  # save locally
-        model.push_to_hub(hf_repo)  # then push to HF repo
-        print(f"✅ Model pushed to HF repo: {hf_repo}")
+       print("📤 Pushing FlashPack model to Hugging Face repo...")
+       logs = push_flashpack_to_hub_local(model, hf_repo)
+       print(logs)
+
 
     return model, dataset, embed_model, tokenizer, long_embeddings