Update app_flash.py

app_flash.py

@@ -1,128 +1,158 @@
-import gc
 import os
+import re
 import torch
 import torch.nn as nn
 import torch.optim as optim
-import gradio as gr
+from typing import Tuple
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM, pipeline
 from flashpack import FlashPackMixin
-from typing import Tuple
+from huggingface_hub import HfApi, create_repo, repo_exists
+import gradio as gr
+from transformers import AutoTokenizer, AutoModel
 
 # ============================================================
-# 🖥 Force CPU mode
+# ⚙️ Setup
 # ============================================================
 device = torch.device("cpu")
 torch.set_num_threads(4)
-print(f"🔧 Forcing device: {device}  (CPU-only mode)")
+print(f"🔧 Using device: {device} (CPU-only mode)")
+
+HF_REPO = "rahul7star/FlashPack"
+MODEL_ID = HF_REPO
+
 
 # ============================================================
-# 1️⃣ Define FlashPack model
+# 🧠 Define FlashPack Trainer
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int = 768, hidden_dim: int = 512, output_dim: int = 768):
+    def __init__(self, input_dim=768, hidden_dim=512, 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: torch.Tensor) -> torch.Tensor:
-        x = self.fc1(x)
-        x = self.relu(x)
-        x = self.fc2(x)
-        return x
+    def forward(self, x):
+        return self.fc2(self.relu(self.fc1(x)))
 
 
 # ============================================================
-# 2️⃣ Build encoder (for embedding)
+# 🔤 Encoder Builder (GPT2 base)
 # ============================================================
-def build_encoder(model_name="gpt2", max_length: int = 32):
+def build_encoder(model_name="gpt2", max_length=32):
     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()
 
     @torch.no_grad()
-    def encode(prompt: str) -> torch.Tensor:
+    def encode(text: str):
         inputs = tokenizer(
-            prompt,
+            text,
             return_tensors="pt",
             truncation=True,
             padding="max_length",
             max_length=max_length,
         ).to(device)
-        outputs = embed_model(**inputs).last_hidden_state.mean(dim=1)
-        return outputs.cpu()
+        return embed_model(**inputs).last_hidden_state.mean(dim=1).cpu()
 
     return tokenizer, embed_model, encode
 
 
 # ============================================================
-# 3️⃣ Load pretrained FlashPack model (skip training)
+# 🧩 FlashPack: Train and Upload (uses Gemma only internally)
 # ============================================================
-def load_flashpack_model(hf_repo="rahul7star/FlashPack"):
-    model = GemmaTrainer.from_flashpack(hf_repo)
-    tokenizer = model.tokenizer if hasattr(model, "tokenizer") else None
-    embed_model = model.embed_model if hasattr(model, "embed_model") else None
-    return model, tokenizer, embed_model
+def train_flashpack_model(hf_repo=HF_REPO):
+    print(f"🚀 Training new FlashPack model for repo: {hf_repo}")
+    model = GemmaTrainer()
+    tokenizer, embed_model, encode = build_encoder("gpt2")
 
-# def load_flashpack_model(hf_repo="rahul7star/FlashPack"):
-#     print(f"🔁 Loading FlashPack model from: {hf_repo}")
-   
-#     model = GemmaTrainer.from_flashpack(hf_repo)
+    # Load dataset (Gemma-expanded dataset)
+    dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
 
-#     model.eval()
-#     tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
-#     return model, tokenizer, embed_model
+    # Compute embeddings for training (short → long)
+    X, Y = [], []
+    for p in dataset.select(range(300)):
+        short_emb = encode(p["short_prompt"])
+        long_emb = encode(p["long_prompt"])
+        X.append(short_emb)
+        Y.append(long_emb)
 
+    X = torch.vstack(X)
+    Y = torch.vstack(Y)
 
-# ============================================================
-# 4️⃣ Load Gemma text model for prompt enhancement
-# ============================================================
-MODEL_ID = "gokaygokay/prompt-enhancer-gemma-3-270m-it"
+    optimizer = optim.Adam(model.parameters(), lr=1e-3)
+    for epoch in range(10):
+        out = model(X)
+        loss = nn.MSELoss()(out, Y)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        print(f"Epoch {epoch+1}/10 | Loss: {loss.item():.6f}")
 
-tokenizer_gemma = AutoTokenizer.from_pretrained(MODEL_ID)
-model_gemma = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+    # Save FlashPack model and push
+    model.to_flashpack("flashpack_model")
+    print("💾 Model saved locally. Uploading to Hugging Face...")
 
-pipe_gemma = pipeline(
-    "text-generation",
-    model=model_gemma,
-    tokenizer=tokenizer_gemma,
-    device=-1,  # CPU
-)
+    api = HfApi()
+    if not repo_exists(hf_repo):
+        create_repo(hf_repo, repo_type="model", exist_ok=True)
+    model.push_to_hub(hf_repo, commit_message="Initial FlashPack model training")
 
-import re
+    print(f"✅ Model uploaded successfully to {hf_repo}")
+    return model, tokenizer, embed_model
 
-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
+
+# ============================================================
+# 📦 Load FlashPack from Hub
+# ============================================================
+def load_flashpack_model(hf_repo=HF_REPO):
+    print(f"📥 Loading FlashPack model from {hf_repo}...")
+    model = GemmaTrainer.from_flashpack(hf_repo)
+    tokenizer, embed_model, encode = build_encoder("gpt2")
+    print("✅ Loaded FlashPack model successfully.")
+    return model, tokenizer, embed_model
+
+
+# ============================================================
+# ⚡ Auto Load or Train
+# ============================================================
+def get_flashpack_model(hf_repo=HF_REPO):
+    try:
+        api = HfApi()
+        if repo_exists(hf_repo):
+            print("✅ Found trained model on Hub.")
+            return load_flashpack_model(hf_repo)
+        else:
+            print("❌ Model not found, training new one using Gemma dataset...")
+            return train_flashpack_model(hf_repo)
+    except Exception as e:
+        print(f"⚠️ Repo check failed: {e}. Retraining model locally.")
+        return train_flashpack_model(hf_repo)
 
 
 # ============================================================
-# 5️⃣ Load FlashPack + Dataset + Encoder
+# 📚 Dataset + Model
 # ============================================================
-model, tokenizer, embed_model = load_flashpack_model("rahul7star/FlashPack")
+model, tokenizer, embed_model = get_flashpack_model()
 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.")
+
+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("✅ FlashPack model and embeddings loaded.")
 
 
 # ============================================================
-# 6️⃣ FlashPack inference helper
+# 🧠 Inference Helpers
 # ============================================================
 @torch.no_grad()
-def encode_for_inference(prompt: str) -> torch.Tensor:
+def encode_for_inference(prompt: str):
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
@@ -140,10 +170,7 @@ def enhance_prompt_flashpack(user_prompt: str, temperature: float, max_tokens: i
     mapped = model(short_emb.to(device)).cpu()
 
     sims = (long_embeddings @ mapped.t()).squeeze(1)
-    long_norms = long_embeddings.norm(dim=1)
-    mapped_norm = mapped.norm()
-    sims = sims / (long_norms * (mapped_norm + 1e-12))
-
+    sims /= (long_embeddings.norm(dim=1) * (mapped.norm() + 1e-12))
     best_idx = int(sims.argmax().item())
     enhanced_prompt = dataset[best_idx]["long_prompt"]
 
@@ -153,36 +180,14 @@ def enhance_prompt_flashpack(user_prompt: str, temperature: float, max_tokens: i
 
 
 # ============================================================
-# 7️⃣ Gemma prompt enhancer
+# 💬 Gradio UI
 # ============================================================
-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
-
-
-# ============================================================
-# 8️⃣ Gradio UI
-# ============================================================
-with gr.Blocks(title="Prompt Enhancer – FlashPack + Gemma (CPU)", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Prompt Enhancer – FlashPack Only", 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.
+    # ✨ FlashPack Prompt Enhancer  
+    - Uses pre-trained **FlashPack model** (`rahul7star/FlashPack`)  
+    - Matches short prompts to enhanced long prompts using learned embeddings  
+    - CPU-only, no Gemma dependency during inference.
     """)
 
     with gr.Row():
@@ -191,17 +196,16 @@ with gr.Blocks(title="Prompt Enhancer – FlashPack + Gemma (CPU)", theme=gr.the
             user_prompt = gr.Textbox(placeholder="Enter a short prompt...", label="Your Prompt", lines=3)
             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)")
+            send_flashpack = gr.Button("🔗 Enhance Prompt")
             clear_btn = gr.Button("🧹 Clear Chat")
 
-    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)
+    user_prompt.submit(enhance_prompt_flashpack, [user_prompt, temperature, max_tokens, chatbot], chatbot)
     clear_btn.click(lambda: [], None, chatbot)
 
+
 # ============================================================
-# 9️⃣ Launch
+# 🚀 Launch App
 # ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)