Update app_flash1.py

app_flash1.py

@@ -9,7 +9,11 @@ from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModel
 from flashpack import FlashPackMixin
 from huggingface_hub import Repository, list_repo_files, hf_hub_download
+import torch.nn.functional as F
 
+# ===========================
+# Device
+# ===========================
 device = torch.device("cpu")
 torch.set_num_threads(4)
 print(f"🔧 Using device: {device} (CPU-only mode)")
@@ -20,7 +24,7 @@ print(f"🔧 Using device: {device} (CPU-only mode)")
 class GemmaTrainer(nn.Module, FlashPackMixin):
     def __init__(self):
         super().__init__()
-        input_dim = 1536  # GPT-2 mean+max pooled embeddings
+        input_dim = 1536
         hidden_dim = 1024
         output_dim = 1536
         self.fc1 = nn.Linear(input_dim, hidden_dim)
@@ -90,7 +94,7 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
 
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
-    # Encode dataset embeddings
+    # Encode embeddings
     s_list, l_list = [], []
     for i, item in enumerate(dataset):
         s_list.append(encode_fn(item["short_prompt"]))
@@ -100,6 +104,9 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
             gc.collect()
     short_emb, long_emb = torch.vstack(s_list), torch.vstack(l_list)
 
+    # Save embeddings & prompts for nearest-neighbor retrieval
+    train_prompts = [item["long_prompt"] for item in dataset]
+
     model = GemmaTrainer()
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
     loss_fn = nn.CosineSimilarity(dim=1)
@@ -118,16 +125,18 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
             break
 
     push_flashpack_model_to_hf(model, hf_repo, log_fn)
-    tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
     @torch.no_grad()
     def enhance_fn(prompt, chat):
         chat = chat or []
-        short_emb = encode_fn(prompt)
-        mapped = model(short_emb.to(device)).cpu()
-        long_prompt = f"🌟 Enhanced prompt (embedding-based) for: {prompt}"
+        short_emb_input = encode_fn(prompt)
+        mapped_emb = model(short_emb_input).cpu()
+        # Nearest neighbor
+        sims = F.cosine_similarity(mapped_emb, long_emb)
+        best_idx = sims.argmax().item()
+        long_prompt = train_prompts[best_idx]
         chat.append({"role": "user", "content": prompt})
-        chat.append({"role": "assistant", "content": long_prompt})
+        chat.append({"role": "assistant", "content": f"🌟 Enhanced prompt: {long_prompt}"})
         return chat
 
     return model, tokenizer, embed_model, enhance_fn, logs
@@ -157,14 +166,20 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     model.eval()
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
+    # Dummy placeholders for nearest neighbor retrieval (replace with actual dataset if available)
+    long_emb = torch.randn(10, 1536)  # placeholder embeddings
+    train_prompts = [f"Example long prompt {i}" for i in range(10)]
+
     @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 (embedding-based) for: {prompt}"
+        short_emb_input = encode_fn(prompt)
+        mapped_emb = model(short_emb_input).cpu()
+        sims = F.cosine_similarity(mapped_emb, long_emb)
+        best_idx = sims.argmax().item()
+        long_prompt = train_prompts[best_idx]
         chat.append({"role": "user", "content": prompt})
-        chat.append({"role": "assistant", "content": long_prompt})
+        chat.append({"role": "assistant", "content": f"🌟 Enhanced prompt: {long_prompt}"})
         return chat
 
     return model, tokenizer, embed_model, enhance_fn