Update app_flash1.py

app_flash1.py

@@ -18,8 +18,11 @@ print(f"🔧 Using device: {device} (CPU-only mode)")
 # Model Definition
 # ===========================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int = 1536, hidden_dim: int = 1024, output_dim: int = 1536):
+    def __init__(self):
         super().__init__()
+        input_dim = 1536  # GPT-2 mean+max pooled embeddings
+        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)
@@ -33,14 +36,14 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         x = self.fc3(x)
         return x
 
-
 # ===========================
 # Encoder
 # ===========================
-def build_encoder(model_name="gpt2", max_length=128):
+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
+
     embed_model = AutoModel.from_pretrained(model_name).to(device)
     embed_model.eval()
 
@@ -48,11 +51,11 @@ def build_encoder(model_name="gpt2", max_length=128):
     def encode(prompt: str) -> torch.Tensor:
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
                            padding="max_length", max_length=max_length).to(device)
-        hidden = embed_model(**inputs).last_hidden_state
-        mean_pool = hidden.mean(dim=1)
-        max_pool, _ = hidden.max(dim=1)
-        return torch.cat([mean_pool, max_pool], dim=1).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()  # doubled embedding
+
     return tokenizer, embed_model, encode
 
 # ===========================
@@ -80,14 +83,14 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
     def log_fn(msg):
         logs.append(msg)
         print(msg)
-    
+
     log_fn("📦 Loading dataset...")
     dataset = load_dataset(dataset_name, split="train").select(range(max_encode))
     log_fn(f"✅ Loaded {len(dataset)} samples")
 
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
-    # Only encode short+long embeddings
+    # Encode dataset embeddings
     s_list, l_list = [], []
     for i, item in enumerate(dataset):
         s_list.append(encode_fn(item["short_prompt"]))
@@ -122,7 +125,7 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
         chat = chat or []
         short_emb = encode_fn(prompt)
         mapped = model(short_emb.to(device)).cpu()
-        long_prompt = f"🌟 Enhanced prompt: {prompt} (creatively expanded)"
+        long_prompt = f"🌟 Enhanced prompt (embedding-based) for: {prompt}"
         chat.append({"role": "user", "content": prompt})
         chat.append({"role": "assistant", "content": long_prompt})
         return chat
@@ -132,7 +135,6 @@ def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
 # ===========================
 # Lazy Load / Get Model
 # ===========================
-# ===========================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     local_model_path = "model.flashpack"
 
@@ -151,8 +153,7 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
             print(f"⚠️ Error accessing HF: {e}")
             return None, None, None, None
 
-    # ⚡ Use input_dim=1536 (default)
-    model = GemmaTrainer(input_dim=1536).from_flashpack(local_model_path)
+    model = GemmaTrainer().from_flashpack(local_model_path)
     model.eval()
     tokenizer, embed_model, encode_fn = build_encoder("gpt2")