Update app_flash.py

app_flash.py

@@ -12,35 +12,17 @@ from huggingface_hub import Repository
 from typing import Tuple
 
 # ============================================================
-# 🖥 Device setup (CPU-only safe)
+# 🖥 Device setup (CPU-only)
 # ============================================================
 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)")
-
+print(f"🔧 Using device: {device} (CPU-only)")
 
 # ============================================================
-# 1️⃣ Define improved FlashPack model
+# 1️⃣ FlashPack model with better hidden layers
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 768):
+    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
         super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
@@ -55,9 +37,8 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         x = self.fc3(x)
         return x
 
-
 # ============================================================
-# 2️⃣ Encoder with mean+max pooling
+# 2️⃣ Encoder using mean+max pooling (for richer embeddings)
 # ============================================================
 def build_encoder(model_name="gpt2", max_length: int = 128):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -69,48 +50,32 @@ def build_encoder(model_name="gpt2", max_length: int = 128):
 
     @torch.no_grad()
     def encode(prompt: str) -> torch.Tensor:
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            padding="max_length",
-            max_length=max_length
-        ).to(device)
+        inputs = tokenizer(prompt, 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)
-        return torch.cat([mean_pool, max_pool], dim=1).cpu()
+        return torch.cat([mean_pool, max_pool], dim=1).cpu()  # doubled embedding
 
     return tokenizer, embed_model, encode
 
-
 # ============================================================
-# 3️⃣ Push FlashPack model to HF
+# 3️⃣ Push FlashPack model to Hugging Face
 # ============================================================
 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"📂 Temporary directory: {tmp_dir}")
         repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
-        logs.append(f"🌐 Hugging Face repo cloned to: {tmp_dir}")
-
         pack_path = os.path.join(tmp_dir, "model.flashpack")
-        logs.append(f"💾 Saving model to: {pack_path}")
         model.save_flashpack(pack_path, target_dtype=torch.float32)
-        logs.append("✅ Model saved successfully.")
-
         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.")
-        logs.append("📄 README.md added.")
-
-        logs.append("🚀 Pushing repo to Hugging Face Hub...")
         repo.push_to_hub()
-        logs.append(f"✅ Model successfully pushed to: {hf_repo}")
-
+        logs.append(f"✅ Model pushed to HF: {hf_repo}")
     return logs
 
-
 # ============================================================
 # 4️⃣ Train FlashPack model
 # ============================================================
@@ -126,9 +91,8 @@ def train_flashpack_model(
     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 (max {max_encode})")
+    print(f"⚡ Using {len(dataset)} prompts for training")
 
-    # Build encoder
     tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
 
     # Encode prompts
@@ -142,14 +106,13 @@ 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]
+    print(f"✅ Encoded embeddings shape: short {short_embeddings.shape}, long {long_embeddings.shape}")
+
+    input_dim = short_embeddings.shape[1]   # should match concatenated mean+max
     output_dim = long_embeddings.shape[1]
 
-    # Build model
     model = GemmaTrainer(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim).to(device)
 
-    # Loss & optimizer
     criterion = nn.CosineSimilarity(dim=1)
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
     max_epochs = 50
@@ -168,7 +131,7 @@ def train_flashpack_model(
 
             optimizer.zero_grad()
             outputs = model(inputs)
-            loss = 1 - criterion(outputs, targets).mean()  # Cosine similarity loss
+            loss = 1 - criterion(outputs, targets).mean()
             loss.backward()
             optimizer.step()
             epoch_loss += loss.item() * inputs.size(0)
@@ -179,8 +142,6 @@ def train_flashpack_model(
 
     print("✅ Training finished!")
 
-    # Push to HF
-    logs = []
     if push_to_hub:
         logs = push_flashpack_model_to_hf(model, hf_repo)
         for log in logs:
@@ -189,31 +150,26 @@ def train_flashpack_model(
     return model, dataset, embed_model, tokenizer, long_embeddings
 
 # ============================================================
-# 5️⃣ Load FlashPack model (train if missing)
+# 5️⃣ Load or train model
 # ============================================================
 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()
-        print("✅ Loaded model successfully from HF")
-        tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
+        tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
         return model, tokenizer, embed_model
     except Exception as e:
         print(f"⚠️ Load failed: {e}")
         print("⏬ Training a new FlashPack model locally...")
         model, dataset, embed_model, tokenizer, long_embeddings = train_flashpack_model()
-        print("📤 Pushing trained model to HF...")
         push_flashpack_model_to_hf(model, hf_repo)
         return model, tokenizer, embed_model, dataset, long_embeddings
 
 # ============================================================
 # 6️⃣ Load or train
 # ============================================================
-try:
-    model, tokenizer, embed_model, dataset, long_embeddings = get_flashpack_model()
-except Exception as e:
-    raise SystemExit(f"❌ Failed to load or train FlashPack model: {e}")
+model, tokenizer, embed_model, dataset, long_embeddings = get_flashpack_model()
 
 # ============================================================
 # 7️⃣ Inference helpers
@@ -221,8 +177,11 @@ except Exception as e:
 @torch.no_grad()
 def encode_for_inference(prompt: str) -> torch.Tensor:
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
-                       padding="max_length", max_length=32).to(device)
-    return embed_model(**inputs).last_hidden_state.mean(dim=1).cpu()
+                       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):
     chat_history = chat_history or []
@@ -268,9 +227,6 @@ with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft
 
 # ============================================================
 # 9️⃣ Launch
-
-# ============================================================
-# 🏁 Launch app
 # ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)