Update app_flash.py

app_flash.py

@@ -67,46 +67,38 @@ 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 = []
+import os
+import tempfile
+from flashpack import FlashPackMixin
+from huggingface_hub import hf_hub_upload
 
-    # 1️⃣ Save model locally to temp directory
-    temp_dir = tempfile.mkdtemp()
-    model.save_flashpack(temp_dir,target_dtype=torch.float32)
-    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 push_flashpack_to_hub_local(model: FlashPackMixin, hf_repo: str):
+    logs = []
 
+    # Create a temporary directory for the pack
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        logs.append(f"📂 Using temporary directory: {tmp_dir}")
+
+        # Ensure a proper file path (not just a directory)
+        pack_file = os.path.join(tmp_dir, "model.pack")
+        logs.append(f"💾 Saving FlashPack model to: {pack_file}")
+
+        # Save the model
+        model.save_flashpack(pack_file, target_dtype=torch.float32)
+        logs.append("✅ Model saved successfully.")
+
+        # Upload to HF Hub
+        logs.append(f"🌐 Uploading {pack_file} to Hugging Face repo: {hf_repo}")
+        hf_hub_upload(
+            repo_id=hf_repo,
+            path_or_fileobj=pack_file,
+            path_in_repo="model.pack",  # filename in the repo
+            repo_type="model",
+            token=os.environ.get("HF_TOKEN")  # make sure HF_TOKEN is set
+        )
+        logs.append("🚀 Upload successful!")
+
+    return logs
 def train_and_push_flashpack(
     dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
     hf_repo: str = "rahul7star/FlashPack",