Update app_flash1.py

app_flash1.py

@@ -9,25 +9,20 @@ 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 pickle
 
-# ============================================================
-# 🖥 Device Setup
-# ============================================================
 device = torch.device("cpu")
 torch.set_num_threads(4)
 print(f"🔧 Using device: {device} (CPU-only mode)")
 
-# ============================================================
-# 1️⃣ Fixed Model Definition (FlashPack-compatible)
-# ============================================================
+# ===========================
+# Model Definition
+# ===========================
 class GemmaTrainer(nn.Module, FlashPackMixin):
     def __init__(self):
         super().__init__()
         input_dim = 1536
         hidden_dim = 1024
         output_dim = 1536
-
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
         self.fc2 = nn.Linear(hidden_dim, hidden_dim)
@@ -41,11 +36,10 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         x = self.fc3(x)
         return x
 
-# ============================================================
-# 2️⃣ Encoder Setup
-# ============================================================
+# ===========================
+# Encoder
+# ===========================
 def build_encoder(model_name="gpt2", max_length=128):
-    print(f"📦 Loading encoder: {model_name}")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -60,54 +54,42 @@ def build_encoder(model_name="gpt2", max_length=128):
         mean_pool = hidden.mean(dim=1)
         max_pool, _ = hidden.max(dim=1)
         return torch.cat([mean_pool, max_pool], dim=1).cpu()
-
+    
     return tokenizer, embed_model, encode
 
-# ============================================================
-# 3️⃣ Push to Hugging Face (model + mapping)
-# ============================================================
-def push_flashpack_model_to_hf(model, short_texts, long_texts, hf_repo):
+# ===========================
+# Push model to HF
+# ===========================
+def push_flashpack_model_to_hf(model, 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 and text mapping pushed to {hf_repo}")
+        print(f"✅ Model pushed to {hf_repo}")
 
-# ============================================================
-# 4️⃣ Training Logic (train + test splits)
-# ============================================================
+# ===========================
+# Training
+# ===========================
 def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
                           hf_repo="rahul7star/FlashPack",
                           max_encode=1000):
     print("📦 Loading dataset...")
-    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")
-
+    dataset = load_dataset(dataset_name, split="train").select(range(max_encode))
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
     def encode_dataset(ds):
-        s_list, l_list, short_texts, long_texts = [], [], [], []
+        s_list, l_list = [], []
         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), short_texts, long_texts
-
-    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)
+        return torch.vstack(s_list), torch.vstack(l_list)
 
+    short_emb, long_emb = encode_dataset(dataset)
     model = GemmaTrainer()
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
     loss_fn = nn.CosineSimilarity(dim=1)
@@ -116,38 +98,40 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
     for epoch in range(20):
         model.train()
         optimizer.zero_grad()
-        preds = model(short_emb_train)
-        loss = 1 - loss_fn(preds, long_emb_train).mean()
+        preds = model(short_emb)
+        loss = 1 - loss_fn(preds, long_emb).mean()
         loss.backward()
         optimizer.step()
-        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}")
-
+        print(f"Epoch {epoch+1}/20 | Loss: {loss.item():.5f}")
         if loss.item() < 0.01:
             print("🎯 Early stopping.")
             break
 
-    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
+    push_flashpack_model_to_hf(model, hf_repo)
+    return model, tokenizer, embed_model
 
-# ============================================================
-# 5️⃣ Load pretrained model for query
-# ============================================================
+# ===========================
+# Load or Train
+# ===========================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
-    print(f"🔍 Checking for model in repo: {hf_repo}")
     local_model_path = "model.flashpack"
 
+    # 1. Try local
     if os.path.exists(local_model_path):
         print("✅ Loading local model")
     else:
-        print("✅ Downloading model from HF")
-        local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
+        # 2. Try HF
+        try:
+            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("🚫 Model not found on HF — will train a new model")
+                return train_flashpack_model(hf_repo=hf_repo)
+        except Exception as e:
+            print(f"⚠️ Error accessing HF: {e}. Training new model instead.")
+            return train_flashpack_model(hf_repo=hf_repo)
 
     model = GemmaTrainer().from_flashpack(local_model_path)
     model.eval()
@@ -158,18 +142,17 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
         chat = chat or []
         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
+        # Simply return a placeholder text for demonstration
+        long_prompt = f"✅ Enhanced long prompt for: {prompt}"
         chat.append({"role": "user", "content": prompt})
         chat.append({"role": "assistant", "content": long_prompt})
         return chat
 
     return model, tokenizer, embed_model, enhance_fn
 
-# ============================================================
-# 6️⃣ Gradio UI
-# ============================================================
+# ===========================
+# Gradio UI
+# ===========================
 with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
     gr.Markdown("## 🧠 FlashPack Prompt Enhancer (CPU)\nShort → Long prompt expander")
 
@@ -180,13 +163,20 @@ with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
     train_btn = gr.Button("🧩 Train Model", variant="secondary")
     status = gr.Markdown("Status: Ready")
 
-    # Load pretrained model
+    # Load or train model
     model, tokenizer, embed_model, enhance_fn = get_flashpack_model()
 
     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(lambda: train_flashpack_model(), None, status)
+
+    def retrain():
+        global model, tokenizer, embed_model, enhance_fn
+        model, tokenizer, embed_model = train_flashpack_model()
+        enhance_fn = get_flashpack_model()[3]
+        return "✅ Model retrained and pushed to HF!"
+
+    train_btn.click(retrain, None, status)
 
 if __name__ == "__main__":
     demo.launch(show_error=True)