Update app_flash.py

app_flash.py

@@ -4,6 +4,7 @@ import torch.optim as optim
 from flashpack import FlashPackMixin
 from datasets import load_dataset
 import gradio as gr
+from transformers import AutoTokenizer, AutoModel
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -28,27 +29,34 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
 # ============================================================
 dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
 
-# Example: convert short_prompt and long_prompt to embeddings
-from transformers import AutoTokenizer, AutoModel
+# ============================================================
+# 3️⃣ Prepare tokenizer & embedding model
+# ============================================================
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token  # FIX padding error
+
 embed_model = AutoModel.from_pretrained("gpt2").to(device)
+embed_model.eval()  # inference only
 
 def encode_prompt(prompt):
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding="max_length", max_length=32).to(device)
     with torch.no_grad():
         return embed_model(**inputs).last_hidden_state.mean(dim=1)
 
-short_embeddings = torch.vstack([encode_prompt(p["short_prompt"]) for p in dataset])
-long_embeddings = torch.vstack([encode_prompt(p["long_prompt"]) for p in dataset])
+# Encode all dataset prompts
+print("📦 Encoding dataset prompts...")
+short_embeddings = torch.vstack([encode_prompt(p["short_prompt"]) for p in dataset]).to(device)
+long_embeddings = torch.vstack([encode_prompt(p["long_prompt"]) for p in dataset]).to(device)
+print(f"✅ Encoded {len(dataset)} prompts")
 
 # ============================================================
-# 3️⃣ Train FlashPack model
+# 4️⃣ Train FlashPack model
 # ============================================================
 model = GemmaTrainer(input_dim=short_embeddings.shape[1], output_dim=long_embeddings.shape[1]).to(device)
 criterion = nn.MSELoss()
 optimizer = optim.Adam(model.parameters(), lr=1e-3)
 
-max_epochs = 1000
+max_epochs = 500
 tolerance = 1e-4
 
 for epoch in range(max_epochs):
@@ -58,57 +66,52 @@ for epoch in range(max_epochs):
     loss.backward()
     optimizer.step()
     if loss.item() < tolerance:
-        print(f"Training converged at epoch {epoch+1}")
+        print(f"✅ Converged at epoch {epoch+1}, Loss={loss.item():.6f}")
         break
-    if epoch % 50 == 0:
-        print(f"Epoch {epoch+1}, Loss: {loss.item():.6f}")
+    if (epoch + 1) % 50 == 0:
+        print(f"Epoch {epoch+1}, Loss={loss.item():.6f}")
 
 # ============================================================
-# 4️⃣ Save to FlashPack Hub
+# 5️⃣ Save FlashPack model to Hub
 # ============================================================
 FLASHPACK_REPO = "rahul7star/FlashPack"
 model.save_flashpack(FLASHPACK_REPO, target_dtype=torch.float32, push_to_hub=True)
-print("✅ Model saved to FlashPack Hub!")
+print(f"✅ Model saved to FlashPack Hub: {FLASHPACK_REPO}")
 
 # ============================================================
-# 5️⃣ Load FlashPack model
+# 6️⃣ Load FlashPack model
 # ============================================================
 loaded_model = model.from_flashpack(FLASHPACK_REPO)
 
 # ============================================================
-# 6️⃣ Gradio interface
+# 7️⃣ Gradio interface
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
-    # Encode short prompt
+    # Encode user prompt
     short_emb = encode_prompt(user_prompt)
-
-    # Generate expanded embedding via trained model
     with torch.no_grad():
         long_emb = loaded_model(short_emb)
 
-    # Decode embedding back to text (approximate via nearest training example)
-    # Simple approach: cosine similarity to long_embeddings
+    # Find nearest matching long prompt in dataset (simple approach)
     cos = nn.CosineSimilarity(dim=1)
     sims = cos(long_emb.repeat(len(long_embeddings),1), long_embeddings)
     best_idx = sims.argmax()
     enhanced_prompt = dataset[best_idx]["long_prompt"]
 
-    # Update chat history
     chat_history.append({"role": "user", "content": user_prompt})
     chat_history.append({"role": "assistant", "content": enhanced_prompt})
     return chat_history
 
 # ============================================================
-# 7️⃣ Gradio UI
+# 8️⃣ Gradio UI
 # ============================================================
 with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         # ✨ Prompt Enhancer (Gemma 3 270M)
-        Enter a short prompt, and the model will **expand it with details and creative context**  
-        using the Gemma chat-template interface.
+        Enter a short prompt, and the model will **expand it with details and creative context**
         """
     )
 
@@ -138,8 +141,5 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
         """
     )
 
-# ============================================================
-# 8️⃣ Launch
-# ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)