Update app_flash.py

testing more encode but a verson before this works just good

app_flash.py

@@ -17,27 +17,49 @@ from typing import Tuple
 device = torch.device("cpu")
 torch.set_num_threads(4)
 print(f"🔧 Using device: {device}  (CPU-only mode)")
+# prompt_enhancer_flashpack_cpu_publish_v2.py
+import gc
+import os
+import tempfile
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModel
+from flashpack import FlashPackMixin
+from huggingface_hub import Repository
+
+device = torch.device("cpu")
+torch.set_num_threads(4)
+print(f"🔧 Using device: {device}  (CPU-only mode)")
+
 
 # ============================================================
-# 1️⃣ Define FlashPack model
+# 1️⃣ Define improved FlashPack model
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int = 768, hidden_dim: int = 512, output_dim: int = 768):
+    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 768):
         super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
-        self.fc2 = nn.Linear(hidden_dim, output_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.fc1(x)
         x = self.relu(x)
         x = self.fc2(x)
+        x = self.relu(x)
+        x = self.fc3(x)
         return x
 
+
 # ============================================================
-# 2️⃣ Build tokenizer + encoder
+# 2️⃣ Encoder with mean+max pooling
 # ============================================================
-def build_encoder(model_name="gpt2", max_length: int = 32):
+def build_encoder(model_name="gpt2", max_length: int = 128):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -54,11 +76,14 @@ 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)
-        return outputs.cpu()
+        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()
 
     return tokenizer, embed_model, encode
 
+
 # ============================================================
 # 3️⃣ Push FlashPack model to HF
 # ============================================================
@@ -85,22 +110,28 @@ def push_flashpack_model_to_hf(model, hf_repo: str):
 
     return logs
 
+
 # ============================================================
 # 4️⃣ Train FlashPack model
 # ============================================================
 def train_flashpack_model(
-    dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
+    dataset_name: str = "rahul7star/prompt-enhancer-dataset",
     max_encode: int = 5000,
-    device: str = "cpu"
+    hidden_dim: int = 1024,
+    push_to_hub: bool = True,
+    hf_repo: str = "rahul7star/FlashPack"
 ) -> 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"⚡ Encoding {len(dataset)} prompts (max {max_encode})")
+    print(f"⚡ Using {len(dataset)} prompts for training (max {max_encode})")
 
-    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
+    # Build encoder
+    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
 
+    # Encode prompts
     short_list, long_list = [], []
     for i, item in enumerate(dataset):
         short_list.append(encode_fn(item["short_prompt"]))
@@ -112,21 +143,20 @@ def train_flashpack_model(
     short_embeddings = torch.vstack(short_list)
     long_embeddings = torch.vstack(long_list)
     print(f"✅ Finished encoding {short_embeddings.shape[0]} prompts")
+    input_dim = short_embeddings.shape[1]
+    output_dim = long_embeddings.shape[1]
 
     # Build model
-    model = GemmaTrainer(
-        input_dim=short_embeddings.shape[1],
-        hidden_dim=min(512, short_embeddings.shape[1]),
-        output_dim=long_embeddings.shape[1]
-    ).to(device)
+    model = GemmaTrainer(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim).to(device)
 
-    criterion = nn.MSELoss()
+    # Loss & optimizer
+    criterion = nn.CosineSimilarity(dim=1)
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
-    max_epochs = 20
+    max_epochs = 30
     batch_size = 32
+    n = short_embeddings.shape[0]
 
     print("🚀 Training model...")
-    n = short_embeddings.shape[0]
     for epoch in range(max_epochs):
         model.train()
         epoch_loss = 0.0
@@ -138,7 +168,7 @@ def train_flashpack_model(
 
             optimizer.zero_grad()
             outputs = model(inputs)
-            loss = criterion(outputs, targets)
+            loss = 1 - criterion(outputs, targets).mean()  # Cosine similarity loss
             loss.backward()
             optimizer.step()
             epoch_loss += loss.item() * inputs.size(0)
@@ -148,6 +178,14 @@ def train_flashpack_model(
             print(f"Epoch {epoch+1}/{max_epochs}, Loss={epoch_loss:.6f}")
 
     print("✅ Training finished!")
+
+    # Push to HF
+    logs = []
+    if push_to_hub:
+        logs = push_flashpack_model_to_hf(model, hf_repo)
+        for log in logs:
+            print(log)
+
     return model, dataset, embed_model, tokenizer, long_embeddings
 
 # ============================================================