Update app_flash.py

app_flash.py

@@ -1,11 +1,11 @@
-# prompt_enhancer_flashpack_cpu_publish.py
 import gc
+import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from datasets import load_dataset
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM, pipeline
 from flashpack import FlashPackMixin
 from typing import Tuple
 
@@ -32,8 +32,9 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         x = self.fc2(x)
         return x
 
+
 # ============================================================
-# 2️⃣ Utility: encode prompts (CPU-friendly)
+# 2️⃣ Build encoder (for embedding)
 # ============================================================
 def build_encoder(model_name="gpt2", max_length: int = 32):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -57,163 +58,60 @@ def build_encoder(model_name="gpt2", max_length: int = 32):
 
     return tokenizer, embed_model, encode
 
-# ============================================================
-# 3️⃣ Train and push FlashPack model
-# ============================================================
-import os
-import tempfile
-from huggingface_hub import hf_hub_download, HfApi
-
-# ------------------------------------------------------------
-# Utility to push FlashPack model to HF using upload_file
-# ------------------------------------------------------------
-import os
-import gc
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import tempfile
-from huggingface_hub import Repository
-from datasets import load_dataset
-from typing import Tuple
-
-# -------------------------------
-# Helper: Push FlashPack model
-# -------------------------------
-def push_flashpack_model_to_hf(model, hf_repo: str):
-    """
-    Save FlashPack model locally and push as Hugging Face model repo.
-    """
-    logs = []
-
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        logs.append(f"📂 Using temporary directory: {tmp_dir}")
-
-        # Clone or initialize HF repo locally
-        repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
-        logs.append(f"🌐 Hugging Face repo cloned to: {tmp_dir}")
-
-        # Save model inside repo
-        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.")
-
-        # Add README
-        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.")
-
-        # Push repo to HF
-        logs.append("🚀 Pushing repo to Hugging Face Hub...")
-        repo.push_to_hub()
-        logs.append(f"✅ Model successfully pushed to: {hf_repo}")
-
-    return logs
-
-# -------------------------------
-# Main training and push function
-# -------------------------------
-def train_and_push_flashpack(
-    dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
-    hf_repo: str = "rahul7star/FlashPack",
-    max_encode: int = 1000,
-    push_to_hub: bool = True,
-    device: str = "cpu"
-) -> Tuple[object, object, object, object, torch.Tensor]:
-
-    print("📦 Loading dataset...")
-    dataset = load_dataset(dataset_name, split="train")
-    limit = min(max_encode, len(dataset))
-    dataset = dataset.select(range(limit))
-    print(f"⚡ Encoding only {len(dataset)} prompts (max limit {max_encode})")
-
-    # Placeholder: build your encoder here
-    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
-
-    print("🔢 Encoding dataset into embeddings (CPU-friendly)...")
-    short_list, long_list = [], []
-    for i, item in enumerate(dataset):
-        short_list.append(encode_fn(item["short_prompt"]))
-        long_list.append(encode_fn(item["long_prompt"]))
-
-        if (i + 1) % 50 == 0 or (i + 1) == len(dataset):
-            print(f"  → Encoded {i+1}/{limit} prompts")
-            gc.collect()
-
-    short_embeddings = torch.vstack(short_list)
-    long_embeddings = torch.vstack(long_list)
-    print(f"✅ Finished encoding {short_embeddings.shape[0]} prompts")
-
-    # Build your FlashPack model (GemmaTrainer placeholder)
-    model = GemmaTrainer(
-        input_dim=short_embeddings.shape[1],
-        hidden_dim=min(512, 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 = 20
-    batch_size = 32
-
-    print("🚀 Training model...")
-    n = short_embeddings.shape[0]
-    for epoch in range(max_epochs):
-        model.train()
-        epoch_loss = 0.0
-        perm = torch.randperm(n)
-        for start in range(0, n, batch_size):
-            idx = perm[start:start+batch_size]
-            inputs = short_embeddings[idx].to(device)
-            targets = long_embeddings[idx].to(device)
-
-            optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
-            loss.backward()
-            optimizer.step()
-            epoch_loss += loss.item() * inputs.size(0)
-
-        epoch_loss /= n
-        if epoch % 5 == 0 or epoch == max_epochs - 1:
-            print(f"Epoch {epoch+1}/{max_epochs}, Loss={epoch_loss:.6f}")
-
-    print("✅ Training finished!")
-
-    logs = []
-    if push_to_hub:
-        print("📤 Pushing model to Hugging Face repo...")
-        logs = push_flashpack_model_to_hf(model, hf_repo)
-        for log in logs:
-            print(log)
-
-    return model, dataset, embed_model, tokenizer, long_embeddings
-
 
 # ============================================================
-# 4️⃣ Load trained model from HF repo
+# 3️⃣ Load pretrained FlashPack model (skip training)
 # ============================================================
 def load_flashpack_model(hf_repo="rahul7star/FlashPack"):
+    print(f"🔁 Loading FlashPack model from: {hf_repo}")
     model = GemmaTrainer.load_flashpack(hf_repo)
     model.eval()
     tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
     return model, tokenizer, embed_model
 
+
 # ============================================================
-# 5️⃣ Run training + push, then reload
+# 4️⃣ Load Gemma text model for prompt enhancement
 # ============================================================
-model, dataset, embed_model, tokenizer, long_embeddings = train_and_push_flashpack(
-    max_encode=1000,  # CPU-safe
-    push_to_hub=True
+MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
+
+tokenizer_gemma = AutoTokenizer.from_pretrained(MODEL_ID)
+model_gemma = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+
+pipe_gemma = pipeline(
+    "text-generation",
+    model=model_gemma,
+    tokenizer=tokenizer_gemma,
+    device=-1,  # CPU
 )
 
-# reload to ensure FlashPack workflow works
+import re
+
+def extract_later_part(user_prompt, generated_text):
+    """Cleans the model output and extracts only the enhanced (later) portion."""
+    cleaned = re.sub(r"<.*?>", "", generated_text).strip()
+    cleaned = re.sub(r"\s+", " ", cleaned)
+    user_prompt_clean = user_prompt.strip().lower()
+    cleaned_lower = cleaned.lower()
+    if cleaned_lower.startswith(user_prompt_clean):
+        cleaned = cleaned[len(user_prompt):].strip(",. ").strip()
+    return cleaned
+
+
+# ============================================================
+# 5️⃣ Load FlashPack + Dataset + Encoder
+# ============================================================
 model, tokenizer, embed_model = load_flashpack_model("rahul7star/FlashPack")
+dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
+long_embeddings = torch.vstack(
+    [embed_model(**tokenizer(p["long_prompt"], return_tensors="pt", truncation=True, padding="max_length", max_length=32)).last_hidden_state.mean(dim=1).cpu()
+     for p in dataset.select(range(min(500, len(dataset))))]
+)
+print("✅ Loaded FlashPack and Gemma models.")
+
 
 # ============================================================
-# 6️⃣ Inference helpers
+# 6️⃣ FlashPack inference helper
 # ============================================================
 @torch.no_grad()
 def encode_for_inference(prompt: str) -> torch.Tensor:
@@ -226,12 +124,13 @@ def encode_for_inference(prompt: str) -> torch.Tensor:
     ).to(device)
     return embed_model(**inputs).last_hidden_state.mean(dim=1).cpu()
 
-def enhance_prompt(user_prompt: str, temperature: float, max_tokens: int, chat_history):
+
+@torch.no_grad()
+def enhance_prompt_flashpack(user_prompt: str, temperature: float, max_tokens: int, chat_history):
     chat_history = chat_history or []
     short_emb = encode_for_inference(user_prompt)
     mapped = model(short_emb.to(device)).cpu()
 
-    cos = nn.CosineSimilarity(dim=1)
     sims = (long_embeddings @ mapped.t()).squeeze(1)
     long_norms = long_embeddings.norm(dim=1)
     mapped_norm = mapped.norm()
@@ -244,33 +143,57 @@ def enhance_prompt(user_prompt: str, temperature: float, max_tokens: int, chat_h
     chat_history.append({"role": "assistant", "content": enhanced_prompt})
     return chat_history
 
+
+# ============================================================
+# 7️⃣ Gemma prompt enhancer
+# ============================================================
+def enhance_prompt_gemma(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}
+    ]
+    prompt = tokenizer_gemma.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    output = pipe_gemma(
+        prompt,
+        max_new_tokens=int(max_tokens),
+        temperature=float(temperature),
+        do_sample=True,
+    )[0]["generated_text"]
+    enhanced_text = extract_later_part(user_prompt, output)
+    chat_history.append({"role": "user", "content": user_prompt})
+    chat_history.append({"role": "assistant", "content": enhanced_text})
+    return chat_history
+
+
 # ============================================================
-# 7️⃣ Gradio UI
+# 8️⃣ Gradio UI
 # ============================================================
-with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # ✨ Prompt Enhancer (FlashPack mapper)
-        Enter a short prompt, and the model will **expand it with details and creative context**.
-        (CPU-only mode.)
-        """
-    )
+with gr.Blocks(title="Prompt Enhancer – FlashPack + Gemma (CPU)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # ✨ Prompt Enhancer (FlashPack + Gemma)
+    - **Gemma model**: Enhances prompts with natural language.
+    - **FlashPack model**: Finds similar expanded prompts from dataset.
+    - CPU-only, for reproducibility.
+    """)
 
     with gr.Row():
-        chatbot = gr.Chatbot(height=400, label="Enhanced Prompts", type="messages")
+        chatbot = gr.Chatbot(height=420, label="Enhanced Prompts", type="messages")
         with gr.Column(scale=1):
             user_prompt = gr.Textbox(placeholder="Enter a short prompt...", label="Your Prompt", lines=3)
-            temperature = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature")
-            max_tokens = gr.Slider(32, 256, value=128, step=16, label="Max Tokens")
-            send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
+            temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
+            max_tokens = gr.Slider(32, 512, value=256, label="Max Tokens")
+            send_gemma = gr.Button("💬 Enhance (Gemma)")
+            send_flashpack = gr.Button("🔗 Enhance (FlashPack)")
             clear_btn = gr.Button("🧹 Clear Chat")
 
-    send_btn.click(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
-    user_prompt.submit(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
+    send_gemma.click(enhance_prompt_gemma, [user_prompt, temperature, max_tokens, chatbot], chatbot)
+    send_flashpack.click(enhance_prompt_flashpack, [user_prompt, temperature, max_tokens, chatbot], chatbot)
+    user_prompt.submit(enhance_prompt_gemma, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
 
 # ============================================================
-# 8️⃣ Launch
+# 9️⃣ Launch
 # ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)