Update app_flash.py

app_flash.py

@@ -153,6 +153,53 @@ def train_flashpack_model(
 # 5️⃣ Load or train model
 # ============================================================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
+    input_dim = 1536  # must match the input_dim used during training
+    try:
+        print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
+
+        # 1️⃣ Try local model first
+        local_model_path = "model.flashpack"
+        if os.path.exists(local_model_path):
+            print("✅ Loading local model")
+        else:
+            # 2️⃣ Try Hugging Face
+            files = list_repo_files(hf_repo)
+            if "model.flashpack" in files:
+                print("✅ Downloading model from HF")
+                from huggingface_hub import hf_hub_download
+                local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
+            else:
+                print("🚫 No pretrained model found")
+                return None, None, None, None
+
+        # 3️⃣ Load model with correct input_dim
+        model = GemmaTrainer(input_dim=input_dim).from_flashpack(local_model_path)
+        model.eval()
+
+        # 4️⃣ Build encoder
+        tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
+
+        # 5️⃣ Enhancement function
+        @torch.no_grad()
+        def enhance_fn(prompt, chat):
+            chat = chat or []
+            short_emb = encode_fn(prompt).to(device)
+            mapped = model(short_emb).cpu()
+            long_prompt = f"🌟 Enhanced prompt: {prompt} (creatively expanded)"
+            chat.append({"role": "user", "content": prompt})
+            chat.append({"role": "assistant", "content": long_prompt})
+            return chat
+
+        return model, tokenizer, embed_model, enhance_fn
+
+    except Exception as e:
+        print(f"⚠️ Load failed: {e}")
+        print("⏬ Training a new FlashPack model locally...")
+        model, dataset, embed_model, tokenizer, long_embeddings = train_flashpack_model()
+        push_flashpack_model_to_hf(model, hf_repo, log_fn=print)
+        return model, tokenizer, embed_model, None
+
+def get_flashpack_model1(hf_repo="rahul7star/FlashPack"):
     try:
         print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
         model = GemmaTrainer.from_flashpack(hf_repo)