Update app_flash.py

app_flash.py

@@ -1,72 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from flashpack import FlashPackMixin
+from datasets import load_dataset
 import gradio as gr
-from transformers import AutoTokenizer
-from flashpack.integrations.transformers import FlashPackTransformersModelMixin
-from transformers import AutoModelForCausalLM, pipeline as hf_pipeline
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ============================================================
-# 1️⃣ Define FlashPack-enabled model class
+# 1️⃣ Define FlashPack model
 # ============================================================
-class FlashPackGemmaModel(AutoModelForCausalLM, FlashPackTransformersModelMixin):
-    """Gemma 3 model wrapped with FlashPackTransformersModelMixin"""
-    pass
+class GemmaTrainer(nn.Module, FlashPackMixin):
+    def __init__(self, input_dim=768, hidden_dim=1024, output_dim=768):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return x
 
 # ============================================================
-# 2️⃣ Load tokenizer
+# 2️⃣ Load dataset
 # ============================================================
-MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
-FLASHPACK_REPO = "rahul7star/FlashPack"
+dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
+
+# Example: convert short_prompt and long_prompt to embeddings
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+embed_model = AutoModel.from_pretrained("gpt2").to(device)
+
+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])
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# ============================================================
+# 3️⃣ 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
+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"Training converged at epoch {epoch+1}")
+        break
+    if epoch % 50 == 0:
+        print(f"Epoch {epoch+1}, Loss: {loss.item():.6f}")
 
 # ============================================================
-# 3️⃣ Load or create FlashPack model
+# 4️⃣ Save to FlashPack Hub
 # ============================================================
-try:
-    print("📂 Loading model from FlashPack repository...")
-    model = FlashPackGemmaModel.from_pretrained_flashpack(FLASHPACK_REPO)
-except FileNotFoundError:
-    print("⚠️ FlashPack model not found. Loading from HF Hub and uploading FlashPack...")
-    model = FlashPackGemmaModel.from_pretrained(MODEL_ID)
-    model.save_pretrained_flashpack(FLASHPACK_REPO, push_to_hub=True)
-    print(f"✅ FlashPack model uploaded to Hugging Face Hub: {FLASHPACK_REPO}")
+FLASHPACK_REPO = "rahul7star/FlashPack"
+model.save_flashpack(FLASHPACK_REPO, target_dtype=torch.float32, push_to_hub=True)
+print("✅ Model saved to FlashPack Hub!")
 
 # ============================================================
-# 4️⃣ Build text-generation pipeline
+# 5️⃣ Load FlashPack model
 # ============================================================
-pipe = hf_pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device_map="auto"
-)
+loaded_model = model.from_flashpack(FLASHPACK_REPO)
 
 # ============================================================
-# 5️⃣ Define prompt enhancement function
+# 6️⃣ Gradio interface
 # ============================================================
 def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
     chat_history = chat_history or []
 
-    messages = [
-        {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
-        {"role": "user", "content": user_prompt},
-    ]
+    # Encode short prompt
+    short_emb = encode_prompt(user_prompt)
 
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    # Generate expanded embedding via trained model
+    with torch.no_grad():
+        long_emb = loaded_model(short_emb)
 
-    outputs = pipe(
-        prompt,
-        max_new_tokens=int(max_tokens),
-        temperature=float(temperature),
-        do_sample=True
-    )
-    enhanced = outputs[0]["generated_text"].strip()
+    # Decode embedding back to text (approximate via nearest training example)
+    # Simple approach: cosine similarity to long_embeddings
+    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})
+    chat_history.append({"role": "assistant", "content": enhanced_prompt})
     return chat_history
 
 # ============================================================
-# 6️⃣ Gradio UI
+# 7️⃣ Gradio UI
 # ============================================================
 with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
@@ -90,7 +125,6 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
             send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
             clear_btn = gr.Button("🧹 Clear Chat")
 
-    # Bind UI actions
     send_btn.click(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     user_prompt.submit(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
@@ -105,7 +139,7 @@ with gr.Blocks(title="Prompt Enhancer – Gemma 3 270M", theme=gr.themes.Soft())
     )
 
 # ============================================================
-# 7️⃣ Launch
+# 8️⃣ Launch
 # ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)