Update app_flash.py

app_flash.py

@@ -7,7 +7,11 @@ from datasets import load_dataset
 import gradio as gr
 from transformers import AutoTokenizer, AutoModel
 
+# ============================================================
+# 🧠 Device setup
+# ============================================================
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🔧 Using device: {device}")
 
 # ============================================================
 # 1️⃣ Define FlashPack model
@@ -25,79 +29,102 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         x = self.fc2(x)
         return x
 
-# ============================================================
-# 2️⃣ Load dataset
-# ============================================================
-dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
 
 # ============================================================
-# 3️⃣ Prepare tokenizer & embedding model
+# 2️⃣ Encode and train using GPU
 # ============================================================
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-tokenizer.pad_token = tokenizer.eos_token  # FIX padding error
+@spaces.GPU(duration=60)  # 10-minute GPU allocation window
+def train_flashpack_model():
+    # Load dataset
+    print("📦 Loading dataset...")
+    dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
 
-embed_model = AutoModel.from_pretrained("gpt2").to(device)
-embed_model.eval()  # inference only
+    # Tokenizer setup
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token  # ✅ Fix padding issue
 
-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)
+    # Base embedding model
+    embed_model = AutoModel.from_pretrained("gpt2").to(device)
+    embed_model.eval()
 
-# 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")
+    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)
 
-# ============================================================
-# 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 = 500
-tolerance = 1e-4
-
-for epoch in range(max_epochs):
-    optimizer.zero_grad()
-    outputs = model(short_embeddings)
-    loss = criterion(outputs, long_embeddings)
-    loss.backward()
-    optimizer.step()
-    if loss.item() < tolerance:
-        print(f"✅ Converged at epoch {epoch+1}, Loss={loss.item():.6f}")
-        break
-    if (epoch + 1) % 50 == 0:
-        print(f"Epoch {epoch+1}, Loss={loss.item():.6f}")
+    # Encode dataset prompts
+    print("🔢 Encoding dataset into embeddings...")
+    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)} pairs")
+
+    # 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 = 500
+    tolerance = 1e-4
+
+    for epoch in range(max_epochs):
+        optimizer.zero_grad()
+        outputs = model(short_embeddings)
+        loss = criterion(outputs, long_embeddings)
+        loss.backward()
+        optimizer.step()
+        if loss.item() < tolerance:
+            print(f"✅ Converged at epoch {epoch+1}, Loss={loss.item():.6f}")
+            break
+        if (epoch + 1) % 50 == 0:
+            print(f"Epoch {epoch+1}, Loss={loss.item():.6f}")
+
+    # Save to Hugging Face Hub
+    FLASHPACK_REPO = "rahul7star/FlashPack"
+    model.save_flashpack(FLASHPACK_REPO, target_dtype=torch.float32, push_to_hub=True)
+    print(f"✅ Model saved to FlashPack Hub: {FLASHPACK_REPO}")
+
+    return model, dataset, embed_model, tokenizer, long_embeddings
 
-# ============================================================
-# 5️⃣ Save FlashPack model to Hub
-# ============================================================
-FLASHPACK_REPO = "rahul7star/FlashPack"
-model.save_flashpack(FLASHPACK_REPO, target_dtype=torch.float32, push_to_hub=True)
-print(f"✅ Model saved to FlashPack Hub: {FLASHPACK_REPO}")
 
 # ============================================================
-# 6️⃣ Load FlashPack model
+# 3️⃣ Run training once and load for inference
 # ============================================================
-loaded_model = model.from_flashpack(FLASHPACK_REPO)
+model, dataset, embed_model, tokenizer, long_embeddings = train_flashpack_model()
+model.eval()
 
 # ============================================================
-# 7️⃣ Gradio interface
+# 4️⃣ Inference function for Gradio
 # ============================================================
+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)
+
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
-
-    # Encode user prompt
     short_emb = encode_prompt(user_prompt)
+
     with torch.no_grad():
-        long_emb = loaded_model(short_emb)
+        long_emb = model(short_emb)
 
-    # Find nearest matching long prompt in dataset (simple approach)
+    # Nearest match search
     cos = nn.CosineSimilarity(dim=1)
-    sims = cos(long_emb.repeat(len(long_embeddings),1), long_embeddings)
+    sims = cos(long_emb.repeat(len(long_embeddings), 1), long_embeddings)
     best_idx = sims.argmax()
     enhanced_prompt = dataset[best_idx]["long_prompt"]
 
@@ -105,8 +132,9 @@ def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history.append({"role": "assistant", "content": enhanced_prompt})
     return chat_history
 
+
 # ============================================================
-# 8️⃣ Gradio UI
+# 5️⃣ Gradio UI
 # ============================================================
 with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft()) as demo:
     gr.Markdown(