Update app_flash1.py

app_flash1.py

@@ -142,41 +142,27 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     print(f"🔍 Checking for model in repo: {hf_repo}")
     local_model_path = "model.flashpack"
-    local_mapping_path = "text_mapping.pkl"
 
-    if os.path.exists(local_model_path) and os.path.exists(local_mapping_path):
-        print("✅ Loading local model and mapping")
+    if os.path.exists(local_model_path):
+        print("✅ Loading local model")
     else:
-        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")
-        if "text_mapping.pkl" in files:
-            print("✅ Downloading text mapping from HF")
-            local_mapping_path = hf_hub_download(repo_id=hf_repo, filename="text_mapping.pkl")
-
-    # Load model
+        print("✅ Downloading model from HF")
+        local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
+
     model = GemmaTrainer().from_flashpack(local_model_path)
     model.eval()
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
-    # Load mapping
-    with open(local_mapping_path, "rb") as f:
-        mapping = pickle.load(f)
-    short_texts, long_texts = mapping["short"], mapping["long"]
-    short_embs = torch.vstack([encode_fn(s) for s in short_texts])
 
-    # Enhance function
     @torch.no_grad()
     def enhance_fn(prompt, chat):
         chat = chat or []
-        query_emb = encode_fn(prompt)
-        mapped = model(query_emb.to(device)).cpu()
-        # Compute cosine similarity to all stored long embeddings
-        sims = torch.nn.functional.cosine_similarity(mapped, short_embs)
-        best_idx = int(sims.argmax())
-        best_long_prompt = long_texts[best_idx]
+        short_emb = encode_fn(prompt)
+        mapped = model(short_emb.to(device)).cpu()
+        # convert mapped tensor into a string (this can be learned in training)
+        # For demonstration, we just return a placeholder
+        long_prompt = f"Enhanced long prompt for: {prompt}"  # replace with your model's actual decoding if available
         chat.append({"role": "user", "content": prompt})
-        chat.append({"role": "assistant", "content": best_long_prompt})
+        chat.append({"role": "assistant", "content": long_prompt})
         return chat
 
     return model, tokenizer, embed_model, enhance_fn