Update app_flash1.py

app_flash1.py

@@ -133,75 +133,40 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
 # Lazy Load / Get Model
 # ===========================
 # ===========================
-# Lazy Load / Get Model (Fixed)
-# ===========================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
-    """
-    Loads the FlashPack model + dataset + long embeddings from HF repo if available,
-    otherwise trains a new model locally.
-    Returns:
-        model, tokenizer, embed_model, enhance_fn, dataset, long_embeddings
-    """
     local_model_path = "model.flashpack"
 
-    try:
-        print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
-        # 1️⃣ Download model from HF
-        files = list_repo_files(hf_repo)
-        if "model.flashpack" in files:
-            local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
-            print("✅ Model downloaded from HF")
-        else:
-            print("🚫 No pretrained model found in HF, will train locally")
-            raise FileNotFoundError
-
-        # 2️⃣ Load FlashPack model
-        model = GemmaTrainer().from_flashpack(local_model_path)
-        model.eval()
-
-        # 3️⃣ Load encoder
-        tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
-
-        # 4️⃣ Try loading dataset + long embeddings from HF
+    if os.path.exists(local_model_path):
+        print("✅ Loading local model")
+    else:
         try:
-            dataset = load_dataset("rahul7star/prompt-enhancer-dataset", split="train").select(range(1000))
-            # Encode long embeddings
-            long_embeddings_list = []
-            for item in dataset:
-                long_embeddings_list.append(encode_fn(item["long_prompt"]))
-            long_embeddings = torch.vstack(long_embeddings_list)
+            files = list_repo_files(hf_repo)
+            if "model.flashpack" in files:
+                print("✅ Downloading model from HF")
+                local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
+            else:
+                print("🚫 No pretrained model found")
+                return None, None, None, None
         except Exception as e:
-            print(f"⚠️ Could not load dataset/embeddings from HF: {e}")
-            dataset = None
-            long_embeddings = None
+            print(f"⚠️ Error accessing HF: {e}")
+            return None, None, None, None
 
-    except Exception:
-        # If anything fails, train locally
-        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)
+    # ⚡ Use input_dim=1536 (default)
+    model = GemmaTrainer(input_dim=1536).from_flashpack(local_model_path)
+    model.eval()
+    tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
-    # 5️⃣ Enhance function using embeddings to select best long prompt
     @torch.no_grad()
     def enhance_fn(prompt, chat):
         chat = chat or []
         short_emb = encode_fn(prompt).to(device)
         mapped = model(short_emb).cpu()
-
-        if dataset is not None and long_embeddings is not None:
-            # Cosine similarity
-            sims = (long_embeddings @ mapped.t()).squeeze(1)
-            sims = sims / (long_embeddings.norm(dim=1) * (mapped.norm() + 1e-12))
-            best_idx = int(sims.argmax().item())
-            enhanced_prompt = dataset[best_idx]["long_prompt"]
-        else:
-            enhanced_prompt = f"🌟 Enhanced prompt (embedding-based) for: {prompt}"
-
+        long_prompt = f"🌟 Enhanced prompt (embedding-based) for: {prompt}"
         chat.append({"role": "user", "content": prompt})
-        chat.append({"role": "assistant", "content": enhanced_prompt})
+        chat.append({"role": "assistant", "content": long_prompt})
         return chat
 
-    return model, tokenizer, embed_model, enhance_fn, dataset, long_embeddings
+    return model, tokenizer, embed_model, enhance_fn
 
 # ===========================
 # Gradio UI