Update app_flash1.py

app_flash1.py

@@ -10,17 +10,19 @@ from transformers import AutoTokenizer, AutoModel
 from flashpack import FlashPackMixin
 from huggingface_hub import Repository
 from typing import Tuple
-from sklearn.model_selection import train_test_split
 
+# ============================================================
+# 🖥 CPU device setup
+# ============================================================
 device = torch.device("cpu")
 torch.set_num_threads(4)
-print(f"🔧 Using device: {device} (CPU-only)")
+print(f"🔧 Using device: {device} (CPU-only mode)")
 
 # ============================================================
-# 1️⃣ Model
+# 1️⃣ FlashPack MLP model (CPU-friendly)
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
+    def __init__(self, input_dim: int, hidden_dim: int = 512, output_dim: int = 768):
         super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
@@ -36,187 +38,160 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         return x
 
 # ============================================================
-# 2️⃣ Encoder with batch mean+max pooling
+# 2️⃣ Lazy-loading GPT-2 encoder
 # ============================================================
-def build_encoder(model_name="gpt2", max_length=128):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    embed_model = AutoModel.from_pretrained(model_name).to(device)
-    embed_model.eval()
+_embed_model = None
+_tokenizer = None
 
-    @torch.no_grad()
-    def encode_batch(prompts: list, batch_size=16) -> torch.Tensor:
-        embeddings = []
-        for i in range(0, len(prompts), batch_size):
-            batch = prompts[i:i+batch_size]
-            inputs = tokenizer(batch, return_tensors="pt", truncation=True,
-                               padding="max_length", max_length=max_length).to(device)
-            last_hidden = embed_model(**inputs).last_hidden_state
-            mean_pool = last_hidden.mean(dim=1)
-            max_pool, _ = last_hidden.max(dim=1)
-            batch_emb = torch.cat([mean_pool, max_pool], dim=1)
-            embeddings.append(batch_emb.cpu())
-        return torch.vstack(embeddings)
+def get_encoder(model_name="gpt2", max_length=64):
+    global _embed_model, _tokenizer
+    if _embed_model is None or _tokenizer is None:
+        print("⚡ Loading GPT-2 encoder model...")
+        _tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if _tokenizer.pad_token is None:
+            _tokenizer.pad_token = _tokenizer.eos_token
+        _embed_model = AutoModel.from_pretrained(model_name).to(device)
+        _embed_model.eval()
+    return _tokenizer, _embed_model
 
-    return tokenizer, embed_model, encode_batch
+@torch.no_grad()
+def encode_prompt(prompt: str) -> torch.Tensor:
+    tokenizer, embed_model = get_encoder()
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
+                       padding="max_length", max_length=64).to(device)
+    last_hidden = embed_model(**inputs).last_hidden_state
+    mean_pool = last_hidden.mean(dim=1)
+    max_pool, _ = last_hidden.max(dim=1)
+    return torch.cat([mean_pool, max_pool], dim=1).cpu()
 
 # ============================================================
-# 3️⃣ Push model to HF
+# 3️⃣ Push FlashPack model to Hugging Face Hub
 # ============================================================
 def push_flashpack_model_to_hf(model, hf_repo: str):
     logs = []
     with tempfile.TemporaryDirectory() as tmp_dir:
-        logs.append(f"📂 Using temporary directory: {tmp_dir}")
+        logs.append(f"📂 Using temp dir: {tmp_dir}")
         repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
         pack_path = os.path.join(tmp_dir, "model.flashpack")
         model.save_flashpack(pack_path, target_dtype=torch.float32)
         readme_path = os.path.join(tmp_dir, "README.md")
         with open(readme_path, "w") as f:
-            f.write("# FlashPack Model\nThis repo contains a FlashPack model trained for short→long prompt mapping.")
+            f.write("# FlashPack Model\nThis repo contains a FlashPack model.")
         repo.push_to_hub()
-        logs.append(f"✅ Model pushed to Hugging Face repo: {hf_repo}")
+        logs.append(f"✅ Model pushed to HF: {hf_repo}")
     return logs
 
 # ============================================================
-# 4️⃣ Train with train/test split & detailed logging
+# 4️⃣ Train FlashPack model with train/test split
 # ============================================================
 def train_flashpack_model(
-    dataset_name="rahul7star/prompt-enhancer-dataset",
-    max_encode=1000,
-    hidden_dim=1024,
-    hf_repo="rahul7star/FlashPack",
-    push_to_hub=True,
-    test_split=0.1,
-    batch_size=32,
-    max_epochs=50,
-    target_test_loss=0.01
+    dataset_name: str = "rahul7star/prompt-enhancer-dataset",
+    max_encode: int = 500,
+    hidden_dim: int = 512,
+    push_to_hub: bool = True,
+    hf_repo: str = "rahul7star/FlashPack",
+    early_stop_threshold: float = 0.001
 ) -> Tuple[GemmaTrainer, object, object, object, torch.Tensor]:
-
+    
     print("📦 Loading dataset...")
     dataset = load_dataset(dataset_name, split="train")
-    limit = min(max_encode, len(dataset))
-    dataset = dataset.select(range(limit))
-    print(f"⚡ Using {len(dataset)} prompts for training")
-
-    short_prompts = [item["short_prompt"] for item in dataset]
-    long_prompts = [item["long_prompt"] for item in dataset]
-
-    # Split
-    train_short, test_short, train_long, test_long = train_test_split(
-        short_prompts, long_prompts, test_size=test_split, random_state=42
-    )
-    print(f"🔹 Train size: {len(train_short)}, Test size: {len(test_short)}")
-
-    tokenizer, embed_model, encode_batch = build_encoder("gpt2", max_length=128)
-
-    # Encode
-    print("⚡ Encoding training prompts...")
-    train_short_emb = encode_batch(train_short)
-    train_long_emb = encode_batch(train_long)
-    print(f"✅ Train embeddings shape: {train_short_emb.shape}, {train_long_emb.shape}")
-
-    print("⚡ Encoding test prompts...")
-    test_short_emb = encode_batch(test_short)
-    test_long_emb = encode_batch(test_long)
-    print(f"✅ Test embeddings shape: {test_short_emb.shape}, {test_long_emb.shape}")
-
-    input_dim = train_short_emb.shape[1]
-    output_dim = train_long_emb.shape[1]
-
+    dataset = dataset.select(range(min(max_encode, len(dataset))))
+    n_train = int(0.8 * len(dataset))
+    n_test = len(dataset) - n_train
+    train_dataset = dataset.select(range(n_train))
+    test_dataset = dataset.select(range(n_train, len(dataset)))
+    print(f"⚡ Train: {n_train}, Test: {n_test}")
+
+    # Encode prompts lazily
+    def batch_encode(ds):
+        short_list, long_list = [], []
+        for i, item in enumerate(ds):
+            short_list.append(encode_prompt(item["short_prompt"]))
+            long_list.append(encode_prompt(item["long_prompt"]))
+            if (i+1) % 20 == 0 or (i+1) == len(ds):
+                print(f"  → Encoded {i+1}/{len(ds)} prompts")
+                gc.collect()
+        return torch.vstack(short_list), torch.vstack(long_list)
+
+    short_train, long_train = batch_encode(train_dataset)
+    short_test, long_test = batch_encode(test_dataset)
+    print(f"✅ Embeddings shapes: short_train={short_train.shape}, long_train={long_train.shape}")
+
+    input_dim = short_train.shape[1]
+    output_dim = long_train.shape[1]
     model = GemmaTrainer(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim).to(device)
 
     criterion = nn.CosineSimilarity(dim=1)
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
-
-    n_train = train_short_emb.shape[0]
+    max_epochs = 50
+    batch_size = 16
 
     print("🚀 Training model...")
+    n = short_train.shape[0]
     for epoch in range(max_epochs):
         model.train()
+        perm = torch.randperm(n)
         epoch_loss = 0.0
-        perm = torch.randperm(n_train)
-        for start in range(0, n_train, batch_size):
+        for start in range(0, n, batch_size):
             idx = perm[start:start+batch_size]
-            inputs = train_short_emb[idx].to(device)
-            targets = train_long_emb[idx].to(device)
-
+            inputs = short_train[idx].to(device)
+            targets = long_train[idx].to(device)
             optimizer.zero_grad()
             outputs = model(inputs)
             loss = 1 - criterion(outputs, targets).mean()
             loss.backward()
             optimizer.step()
             epoch_loss += loss.item() * inputs.size(0)
+        epoch_loss /= n
 
-        epoch_loss /= n_train
-
-        # Evaluate on test
+        # Evaluate on test set
         model.eval()
         with torch.no_grad():
-            test_outputs = model(test_short_emb.to(device))
-            test_loss = (1 - criterion(test_outputs, test_long_emb.to(device)).mean()).item()
+            outputs_test = model(short_test.to(device))
+            test_loss = 1 - criterion(outputs_test, long_test.to(device)).mean().item()
 
-        print(f"Epoch {epoch+1}/{max_epochs} → Train loss: {epoch_loss:.6f}, Test loss: {test_loss:.6f}")
+        print(f"Epoch {epoch+1}/{max_epochs} | Train Loss={epoch_loss:.6f} | Test Loss={test_loss:.6f}")
 
-        # Check if model is perfect enough
-        if test_loss <= target_test_loss:
-            print(f"✅ Target test loss reached ({test_loss:.6f}) – stopping training early.")
+        # Early stop: very low test loss means model is good
+        if test_loss < early_stop_threshold:
+            print("🎯 Early stop: test loss below threshold. Model is ready!")
             break
 
-    # Push to HF if trained well
-    logs = []
-    if push_to_hub and test_loss <= target_test_loss:
+    if push_to_hub:
         logs = push_flashpack_model_to_hf(model, hf_repo)
         for log in logs:
             print(log)
-    elif push_to_hub:
-        print(f"⚠️ Test loss too high ({test_loss:.6f}); skipping HF upload.")
 
-    return model, dataset, embed_model, tokenizer, train_long_emb
+    return model, dataset, None, None, long_train  # embed_model and tokenizer lazy-loaded
 
 # ============================================================
-# 5️⃣ Load or train
+# 5️⃣ Lazy load or train
 # ============================================================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     try:
         print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
         model = GemmaTrainer.from_flashpack(hf_repo)
         model.eval()
-        tokenizer, embed_model, encode_batch = build_encoder("gpt2", max_length=128)
-        return model, tokenizer, embed_model
+        print("✅ Loaded model from HF")
+        return model
     except Exception as e:
         print(f"⚠️ Load failed: {e}")
-        print("⏬ Training a new FlashPack model locally...")
-        return train_flashpack_model(hf_repo=hf_repo)
+        print("⏬ Training new FlashPack model locally...")
+        model, dataset, _, _, long_embeddings = train_flashpack_model()
+        return model, dataset, long_embeddings
 
 # ============================================================
-# 6️⃣ Load or train
-# ============================================================
-model, tokenizer, embed_model, dataset, long_embeddings = get_flashpack_model()
-
-# ============================================================
-# 7️⃣ Inference helpers
+# 6️⃣ Inference helpers
 # ============================================================
 @torch.no_grad()
-def encode_for_inference(prompt: str) -> torch.Tensor:
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
-                       padding="max_length", max_length=128).to(device)
-    last_hidden = embed_model(**inputs).last_hidden_state
-    mean_pool = last_hidden.mean(dim=1)
-    max_pool, _ = last_hidden.max(dim=1)
-    return torch.cat([mean_pool, max_pool], dim=1).cpu()
-
-def enhance_prompt(user_prompt: str, temperature: float, max_tokens: int, chat_history):
+def enhance_prompt(user_prompt: str, chat_history, model, long_embeddings, dataset):
     chat_history = chat_history or []
-    short_emb = encode_for_inference(user_prompt)
+    short_emb = encode_prompt(user_prompt)
     mapped = model(short_emb.to(device)).cpu()
-
     sims = (long_embeddings @ mapped.t()).squeeze(1)
     long_norms = long_embeddings.norm(dim=1)
     mapped_norm = mapped.norm()
     sims = sims / (long_norms * (mapped_norm + 1e-12))
-
     best_idx = int(sims.argmax().item())
     enhanced_prompt = dataset[best_idx]["long_prompt"]
 
@@ -225,28 +200,21 @@ def enhance_prompt(user_prompt: str, temperature: float, max_tokens: int, chat_h
     return chat_history
 
 # ============================================================
-# 8️⃣ Gradio UI
+# 7️⃣ Launch Gradio app
 # ============================================================
-with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # ✨ Prompt Enhancer (FlashPack mapper)
-        Enter a short prompt, and the model will **expand it with details and creative context**.
-        (CPU-only mode.)
-        """
-    )
+model, dataset, long_embeddings = get_flashpack_model()
 
+with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ✨ Prompt Enhancer (FlashPack mapper)\nEnter a short prompt, and it will expand it.")
     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)
-            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")
             clear_btn = gr.Button("🧹 Clear Chat")
 
-    send_btn.click(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
-    user_prompt.submit(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
+    send_btn.click(enhance_prompt, [user_prompt, chatbot, model, long_embeddings, dataset], chatbot)
+    user_prompt.submit(enhance_prompt, [user_prompt, chatbot, model, long_embeddings, dataset], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":