Update app_flash.py

app_flash.py

@@ -6,14 +6,14 @@ import torch.optim as optim
 from datasets import load_dataset
 import gradio as gr
 from transformers import AutoTokenizer, AutoModel
-from flashpack import FlashPackMixin  # keep if your mixin provides save_flashpack
+from flashpack import FlashPackMixin
 from typing import Tuple
 
 # ============================================================
-# 🖥 Force CPU mode (safe for HF Spaces / Kaggle)
+# 🖥 Force CPU mode
 # ============================================================
 device = torch.device("cpu")
-torch.set_num_threads(4)  # reduce CPU contention in shared environments
+torch.set_num_threads(4)  # reduce CPU contention
 print(f"🔧 Forcing device: {device}  (CPU-only mode)")
 
 # ============================================================
@@ -37,7 +37,6 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
 # ============================================================
 def build_encoder(model_name="gpt2", max_length: int = 32):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    # Some GPT2 tokenizers have no pad token — set eos as pad
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
@@ -46,10 +45,6 @@ def build_encoder(model_name="gpt2", max_length: int = 32):
 
     @torch.no_grad()
     def encode(prompt: str) -> torch.Tensor:
-        """
-        Encodes a single prompt and returns a CPU tensor of shape (1, hidden_size).
-        Always returns a CPU tensor to avoid device juggling in downstream code.
-        """
         inputs = tokenizer(
             prompt,
             return_tensors="pt",
@@ -57,8 +52,7 @@ def build_encoder(model_name="gpt2", max_length: int = 32):
             padding="max_length",
             max_length=max_length,
         ).to(device)
-
-        outputs = embed_model(**inputs).last_hidden_state.mean(dim=1)  # (1, hidden)
+        outputs = embed_model(**inputs).last_hidden_state.mean(dim=1)
         return outputs.cpu()
 
     return tokenizer, embed_model, encode
@@ -70,10 +64,10 @@ def train_flashpack_model(
     dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
     model_name: str = "gpt2",
     max_length: int = 32,
-    max_encode: int = 2000,  # maximum number of prompts to encode
+    max_encode: int = 1000,  # use smaller number for CPU
     push_to_hub: bool = False,
     hf_repo: str = "rahul7star/FlashPack",
-) -> tuple:
+) -> Tuple[GemmaTrainer, object, object, object, torch.Tensor]:
 
     # 1️⃣ Load dataset
     print("📦 Loading dataset...")
@@ -84,23 +78,17 @@ def train_flashpack_model(
     dataset = dataset.select(range(limit))
     print(f"⚡ Encoding only {len(dataset)} prompts (max limit {max_encode})")
 
-    # 2️⃣ Setup tokenizer & encoder
-    tokenizer, embed_model, encode_fn = build_encoder(model_name=model_name, max_length=max_length)
+    # 2️⃣ Setup encoder
+    tokenizer, embed_model, encode_fn = build_encoder(model_name, max_length)
 
-    # 3️⃣ Encode dataset (CPU-friendly)
+    # 3️⃣ Encode dataset
     print("🔢 Encoding dataset into embeddings (CPU-friendly)...")
     short_list, long_list = [], []
     for i, item in enumerate(dataset):
         short_list.append(encode_fn(item["short_prompt"]))
         long_list.append(encode_fn(item["long_prompt"]))
 
-        # Exit early if we hit max_encode
-        if (i + 1) >= max_encode:
-            print(f"⚡ Reached max encode limit: {max_encode} prompts, stopping early.")
-            break
-
-        # Progress logging
-        if (i + 1) % 50 == 0:
+        if (i + 1) % 50 == 0 or (i + 1) == len(dataset):
             print(f"  → Encoded {i+1}/{limit} prompts")
             gc.collect()
 
@@ -108,7 +96,7 @@ def train_flashpack_model(
     long_embeddings = torch.vstack(long_list)
     print(f"✅ Finished encoding {short_embeddings.shape[0]} prompts")
 
-    # 4️⃣ Initialize and train model (same as before)
+    # 4️⃣ Initialize & train model
     model = GemmaTrainer(
         input_dim=short_embeddings.shape[1],
         hidden_dim=min(512, short_embeddings.shape[1]),
@@ -117,8 +105,7 @@ def train_flashpack_model(
 
     criterion = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
-    max_epochs = 50
-    tolerance = 1e-4
+    max_epochs = 20
     batch_size = 32
 
     print("🚀 Training FlashPack mapper model (CPU)...")
@@ -143,26 +130,28 @@ def train_flashpack_model(
         if epoch % 5 == 0 or epoch == max_epochs-1:
             print(f"Epoch {epoch+1}/{max_epochs}, Loss={epoch_loss:.6f}")
 
-        if epoch_loss < tolerance:
-            print(f"✅ Converged at epoch {epoch+1}, Loss={epoch_loss:.6f}")
-            break
-
     print("✅ Training finished!")
+
+    # 5️⃣ Push to HF repo if requested
+    if push_to_hub:
+        model.save_flashpack(hf_repo, target_dtype=torch.float32, push_to_hub=True)
+        print(f"✅ Model pushed to HF repo: {hf_repo}")
+
     return model, dataset, embed_model, tokenizer, long_embeddings
 
 # ============================================================
-# 4️⃣ Build everything and prepare for inference
+# 4️⃣ Run training & load model
 # ============================================================
-# For demo speed in CPU mode, you might want a subset_limit (e.g., 1000).
-# Set subset_limit=None to use full dataset.
 model, dataset, embed_model, tokenizer, long_embeddings = train_flashpack_model(
-    subset_limit=None,  # change to a small int for faster testing
-    push_to_hub=False,  # toggle when you want to actually push
+    max_encode=1000,  # safe CPU-friendly subset
+    push_to_hub=False
 )
 
 model.eval()
 
-# Reusable encode function for inference (returns CPU tensor)
+# ============================================================
+# 5️⃣ Inference helpers
+# ============================================================
 @torch.no_grad()
 def encode_for_inference(prompt: str) -> torch.Tensor:
     inputs = tokenizer(
@@ -174,22 +163,13 @@ def encode_for_inference(prompt: str) -> torch.Tensor:
     ).to(device)
     return embed_model(**inputs).last_hidden_state.mean(dim=1).cpu()
 
-# ============================================================
-# 5️⃣ Enhance prompt function (nearest neighbor via cosine)
-# ============================================================
 def enhance_prompt(user_prompt: str, temperature: float, max_tokens: int, chat_history):
     chat_history = chat_history or []
+    short_emb = encode_for_inference(user_prompt)
+    mapped = model(short_emb.to(device)).cpu()
 
-    # encode user prompt (CPU tensor)
-    short_emb = encode_for_inference(user_prompt)  # (1, dim)
-    with torch.no_grad():
-        mapped = model(short_emb.to(device)).cpu()   # (1, dim)
-
-    # cosine similarity against dataset long embeddings
     cos = nn.CosineSimilarity(dim=1)
-    # mapped.repeat(len(long_embeddings), 1) is heavy; do efficient matmul similarity:
     sims = (long_embeddings @ mapped.t()).squeeze(1)
-    # normalize: sims / (||long|| * ||mapped||)
     long_norms = long_embeddings.norm(dim=1)
     mapped_norm = mapped.norm()
     sims = sims / (long_norms * (mapped_norm + 1e-12))
@@ -209,18 +189,14 @@ with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft
         """
         # ✨ Prompt Enhancer (FlashPack mapper)
         Enter a short prompt, and the model will **expand it with details and creative context**.
-        (This demo runs on CPU — expect slower inference/training than GPU.)
+        (CPU-only mode.)
         """
     )
 
     with gr.Row():
         chatbot = gr.Chatbot(height=400, label="Enhanced Prompts", type="messages")
         with gr.Column(scale=1):
-            user_prompt = gr.Textbox(
-                placeholder="Enter a short prompt...",
-                label="Your Prompt",
-                lines=3,
-            )
+            user_prompt = gr.Textbox(placeholder="Enter a short prompt...", label="Your Prompt", lines=3)
             temperature = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature")
             max_tokens = gr.Slider(32, 256, value=128, step=16, label="Max Tokens")
             send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
@@ -230,15 +206,6 @@ with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft
     user_prompt.submit(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
 
-    gr.Markdown(
-        """
-        ---
-        💡 **Tips:**
-        - CPU mode: training and large-batch encodes can take a while. Use `subset_limit` in the training call for quick tests.
-        - Increase *Temperature* for more creative outputs (not used in the nearest-neighbour mapper but kept for UI parity).
-        """
-    )
-
 # ============================================================
 # 7️⃣ Launch
 # ============================================================