Update app_flash.py

app_flash.py

@@ -106,12 +106,28 @@ def train_flashpack_model(hf_repo=HF_REPO):
 # ============================================================
 # 📦 Load FlashPack from Hub
 # ============================================================
-def load_flashpack_model(hf_repo=HF_REPO):
-    print(f"📥 Loading FlashPack model from {hf_repo}...")
-    model = GemmaTrainer.from_flashpack(hf_repo)
-    tokenizer, embed_model, encode = build_encoder("gpt2")
-    print("✅ Loaded FlashPack model successfully.")
-    return model, tokenizer, embed_model
+from huggingface_hub import snapshot_download
+import os
+
+def load_flashpack_model(hf_repo="rahul7star/FlashPack"):
+    print(f"🔁 Loading FlashPack model from: {hf_repo}")
+
+    # Try local first, then Hugging Face Hub
+    if os.path.isdir(hf_repo):
+        local_dir = hf_repo
+        print(f"📂 Using local FlashPack model at: {local_dir}")
+    else:
+        print("☁️ Downloading FlashPack model from Hugging Face Hub...")
+        local_dir = snapshot_download(repo_id=hf_repo)
+        print(f"📥 Model snapshot downloaded to: {local_dir}")
+
+    # Load from local directory
+    model = GemmaTrainer.from_flashpack(local_dir)
+    model.eval()
+    print("✅ FlashPack model loaded successfully.")
+    return model
+
+
 
 
 # ============================================================