Update app_flash1.py

app_flash1.py

@@ -9,7 +9,7 @@ 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
-from typing import Tuple
+import pickle
 
 # ============================================================
 # 🖥 Device Setup
@@ -64,43 +64,51 @@ def build_encoder(model_name="gpt2", max_length=128):
     return tokenizer, embed_model, encode
 
 # ============================================================
-# 3️⃣ Push to Hugging Face
+# 3️⃣ Push to Hugging Face (model + mapping)
 # ============================================================
-def push_flashpack_model_to_hf(model, hf_repo):
+def push_flashpack_model_to_hf(model, short_texts, long_texts, hf_repo):
     with tempfile.TemporaryDirectory() as tmp_dir:
         repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
+        # Save model
         model.save_flashpack(os.path.join(tmp_dir, "model.flashpack"))
+        # Save text mapping
+        with open(os.path.join(tmp_dir, "text_mapping.pkl"), "wb") as f:
+            pickle.dump({"short": short_texts, "long": long_texts}, f)
+        # README
         with open(os.path.join(tmp_dir, "README.md"), "w") as f:
             f.write("# FlashPack Model\nTrained locally and pushed to HF.")
         repo.push_to_hub()
-        print(f"✅ Model pushed to {hf_repo}")
+        print(f"✅ Model and text mapping pushed to {hf_repo}")
 
 # ============================================================
-# 4️⃣ Training Logic
+# 4️⃣ Training Logic (train + test splits)
 # ============================================================
 def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
                           hf_repo="rahul7star/FlashPack",
                           max_encode=1000):
     print("📦 Loading dataset...")
-    dataset = load_dataset(dataset_name, split="train").select(range(max_encode))
-    print(f"✅ Loaded {len(dataset)} samples")
+    dataset_train = load_dataset(dataset_name, split="train").select(range(max_encode))
+    dataset_test = load_dataset(dataset_name, split="test").select(range(max_encode // 10))
+    print(f"✅ Loaded {len(dataset_train)} train and {len(dataset_test)} test samples")
 
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
     def encode_dataset(ds):
-        s_list, l_list = [], []
+        s_list, l_list, short_texts, long_texts = [], [], [], []
         for i, item in enumerate(ds):
             s_list.append(encode_fn(item["short_prompt"]))
             l_list.append(encode_fn(item["long_prompt"]))
+            short_texts.append(item["short_prompt"])
+            long_texts.append(item["long_prompt"])
             if (i + 1) % 50 == 0:
                 print(f"  → Encoded {i + 1}/{len(ds)}")
                 gc.collect()
-        return torch.vstack(s_list), torch.vstack(l_list)
+        return torch.vstack(s_list), torch.vstack(l_list), short_texts, long_texts
 
-    short_emb, long_emb = encode_dataset(dataset)
-    input_dim, output_dim = short_emb.shape[1], long_emb.shape[1]
-    model = GemmaTrainer(input_dim, 1024, output_dim)
+    short_emb_train, long_emb_train, short_texts_train, long_texts_train = encode_dataset(dataset_train)
+    short_emb_test, long_emb_test, _, _ = encode_dataset(dataset_test)
 
+    model = GemmaTrainer()
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
     loss_fn = nn.CosineSimilarity(dim=1)
 
@@ -108,122 +116,91 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
     for epoch in range(20):
         model.train()
         optimizer.zero_grad()
-        preds = model(short_emb)
-        loss = 1 - loss_fn(preds, long_emb).mean()
+        preds = model(short_emb_train)
+        loss = 1 - loss_fn(preds, long_emb_train).mean()
         loss.backward()
         optimizer.step()
-        print(f"Epoch {epoch+1}/20 | Loss: {loss.item():.5f}")
+        print(f"Epoch {epoch+1}/20 | Train Loss: {loss.item():.5f}")
+
+        # Evaluate on test
+        model.eval()
+        with torch.no_grad():
+            test_preds = model(short_emb_test)
+            test_loss = 1 - loss_fn(test_preds, long_emb_test).mean()
+        print(f"           | Test Loss: {test_loss.item():.5f}")
+
         if loss.item() < 0.01:
             print("🎯 Early stopping.")
             break
 
-    push_flashpack_model_to_hf(model, hf_repo)
-    return model, tokenizer, embed_model, dataset, long_emb
+    push_flashpack_model_to_hf(model, short_texts_train, long_texts_train, hf_repo)
+    return model, tokenizer, embed_model, short_emb_train, long_emb_train, short_texts_train, long_texts_train
 
 # ============================================================
-# 5️⃣ Load or Train
+# 5️⃣ Load pretrained model for query
 # ============================================================
 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"
 
-    local_path = "model.flashpack"
-
-    # 1️⃣ Try local first
-    if os.path.exists(local_path):
-        print("✅ Found local model.flashpack — loading it directly.")
-        model = GemmaTrainer().from_flashpack(local_path)
-        model.eval()
-        tokenizer, embed_model, _ = build_encoder("gpt2")
+    if os.path.exists(local_model_path) and os.path.exists(local_mapping_path):
+        print("✅ Loading local model and mapping")
     else:
-        # 2️⃣ Check HF repo
-        try:
-            files = list_repo_files(hf_repo)
-            if "model.flashpack" in files:
-                print("✅ Found model.flashpack in repo — downloading and loading it.")
-                local_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
-                model = GemmaTrainer().from_flashpack(local_path)
-                model.eval()
-                tokenizer, embed_model, _ = build_encoder("gpt2")
-            else:
-                print("🚫 model.flashpack not found — starting training.")
-                return train_flashpack_model(hf_repo=hf_repo)
-        except Exception as e:
-            print(f"⚠️ Error checking repo: {e}")
-            print("⏬ Training new model instead.")
-            return train_flashpack_model(hf_repo=hf_repo)
-
-    # ✅ Enhance function without dataset
-    def enhance_fn(prompt, chat):
-        chat = chat or []
-        short_emb = encode_prompt(prompt, tokenizer, embed_model)
-        mapped = model(short_emb.to(device)).cpu()
-        # We don't need a dataset; just return the mapped tensor info as string
-        chat.append({"role": "user", "content": prompt})
-        chat.append({"role": "assistant", "content": f"✅ Model loaded — ready to enhance.\nOutput vector: {mapped[0].tolist()[:8]} ..."})
-        return chat
-
-    return model, tokenizer, embed_model, None, None, enhance_fn
-
-
-
-# ============================================================
-# 6️⃣ Encode & Enhance Functions
-# ============================================================
-@torch.no_grad()
-def encode_prompt(prompt, tokenizer, embed_model):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
-                       padding="max_length", max_length=128).to(device)
-    hidden = embed_model(**inputs).last_hidden_state
-    mean_pool = hidden.mean(dim=1)
-    max_pool, _ = hidden.max(dim=1)
-    return torch.cat([mean_pool, max_pool], dim=1).cpu()
-
+        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
+    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])
 
-def make_enhance_fn(model, tokenizer, embed_model, long_emb, dataset):
+    # Enhance function
     @torch.no_grad()
-    def fn(prompt, chat):
+    def enhance_fn(prompt, chat):
         chat = chat or []
-        short_emb = encode_prompt(prompt, tokenizer, embed_model)
-        mapped = model(short_emb.to(device)).cpu()
-        sims = (long_emb @ mapped.t()).squeeze(1)
-        best = int(sims.argmax())
-        enhanced = dataset[best]["long_prompt"]
+        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]
         chat.append({"role": "user", "content": prompt})
-        chat.append({"role": "assistant", "content": enhanced})
+        chat.append({"role": "assistant", "content": best_long_prompt})
         return chat
-    return fn
+
+    return model, tokenizer, embed_model, enhance_fn
 
 # ============================================================
-# 7️⃣ Gradio UI
+# 6️⃣ Gradio UI
 # ============================================================
 with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
     gr.Markdown("## 🧠 FlashPack Prompt Enhancer (CPU)\nShort → Long prompt expander")
 
     chatbot = gr.Chatbot(height=400, type="messages")
-
     user_input = gr.Textbox(label="Your prompt")
     send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
     clear_btn = gr.Button("🧹 Clear")
     train_btn = gr.Button("🧩 Train Model", variant="secondary")
-
     status = gr.Markdown("Status: Ready")
 
-    # Load model initially
-    model, tokenizer, embed_model, dataset, long_emb, enhance_fn = get_flashpack_model()
-
-    def enhance(prompt, chat):
-        return enhance_fn(prompt, chat)
-
-    def retrain():
-        global model, tokenizer, embed_model, dataset, long_emb, enhance_fn
-        model, tokenizer, embed_model, dataset, long_emb = train_flashpack_model()
-        enhance_fn = make_enhance_fn(model, tokenizer, embed_model, long_emb, dataset)
-        return "✅ Model retrained and pushed to HF!"
+    # Load pretrained model
+    model, tokenizer, embed_model, enhance_fn = get_flashpack_model()
 
-    send_btn.click(enhance, [user_input, chatbot], chatbot)
-    user_input.submit(enhance, [user_input, chatbot], chatbot)
+    send_btn.click(enhance_fn, [user_input, chatbot], chatbot)
+    user_input.submit(enhance_fn, [user_input, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
-    train_btn.click(retrain, None, status)
+    train_btn.click(lambda: train_flashpack_model(), None, status)
 
 if __name__ == "__main__":
     demo.launch(show_error=True)