Update app_flash1.py

app_flash1.py

@@ -1,139 +1,232 @@
-import gradio as gr
+import gc
+import os
 import torch
+import torch.nn as nn
+import torch.optim as optim
+import tempfile
+import gradio as gr
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModel
+from flashpack import FlashPackMixin
 from huggingface_hub import Repository
-import tempfile, os, gc, torch.nn as nn, torch.optim as optim
+from typing import Tuple
+
+# ============================================================
+# 🖥 Device setup (CPU-only)
+# ============================================================
+device = torch.device("cpu")
+torch.set_num_threads(4)
+print(f"🔧 Using device: {device} (CPU-only mode)")
 
 # ============================================================
-# GemmaTrainer simple model
+# 1️⃣ FlashPack model
 # ============================================================
-class GemmaTrainer(nn.Module):
-    def __init__(self, input_dim, hidden_dim, output_dim):
+class GemmaTrainer(nn.Module, FlashPackMixin):
+    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
         super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
-        self.fc2 = nn.Linear(hidden_dim, output_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, output_dim)
 
-    def forward(self, x):
-        return self.fc2(self.relu(self.fc1(x)))
-
-    def save_flashpack(self, path, target_dtype=torch.float32):
-        torch.save(self.state_dict(), path)
-
-    @classmethod
-    def from_flashpack(cls, repo_path, model=None):
-        local_path = os.path.expanduser(
-            f"~/.cache/huggingface/hub/models--{repo_path.replace('/', '--')}/snapshots/"
-        )
-        # Find the newest snapshot
-        for root, dirs, files in os.walk(local_path):
-            if "model.flashpack" in files:
-                file_path = os.path.join(root, "model.flashpack")
-                model.load_state_dict(torch.load(file_path, map_location="cpu"))
-                return model
-        raise FileNotFoundError("model.flashpack not found in repo cache")
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.relu(x)
+        x = self.fc3(x)
+        return x
 
 # ============================================================
-# Build encoder (tokenizer + model)
+# 2️⃣ Encoder (mean + max pooling)
 # ============================================================
-def build_encoder(model_name="gpt2", max_length: int = 32):
+def build_encoder(model_name="gpt2", max_length: int = 128):
+    print(f"📦 Loading tokenizer and model for {model_name}...")
     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)
+
+    embed_model = AutoModel.from_pretrained(model_name).to(device)
     embed_model.eval()
+    print(f"✅ Encoder ready: {model_name}")
 
     @torch.no_grad()
     def encode(prompt: str) -> torch.Tensor:
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            padding="max_length",
-            max_length=max_length
-        )
-        outputs = embed_model(**inputs).last_hidden_state.mean(dim=1)
-        return outputs.cpu()
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
+                           padding="max_length", max_length=max_length).to(device)
+        last_hidden = embed_model(**inputs).last_hidden_state
+        mean_pool = last_hidden.mean(dim=1)
+        max_pool, _ = last_hidden.max(dim=1)
+        return torch.cat([mean_pool, max_pool], dim=1).cpu()  # double dimension
 
     return tokenizer, embed_model, encode
 
 # ============================================================
-# Train FlashPack model (only if not found)
+# 3️⃣ Push FlashPack model to HF
+# ============================================================
+def push_flashpack_model_to_hf(model, hf_repo: str):
+    logs = []
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        logs.append(f"📂 Temporary directory: {tmp_dir}")
+        repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
+        pack_path = os.path.join(tmp_dir, "model.flashpack")
+        model.save_flashpack(pack_path, target_dtype=torch.float32)
+        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.")
+        repo.push_to_hub()
+        logs.append(f"✅ Model pushed to HF: {hf_repo}")
+    return logs
+
+# ============================================================
+# 4️⃣ Train FlashPack model
 # ============================================================
 def train_flashpack_model(
     dataset_name: str = "rahul7star/prompt-enhancer-dataset",
-    max_encode: int = 500,
-    device: str = "cpu"
-):
+    max_encode: int = 1000,
+    hidden_dim: int = 1024,
+    hf_repo: str = "rahul7star/FlashPack",
+    push_to_hub: bool = True
+) -> Tuple[GemmaTrainer, object, object, object, torch.Tensor]:
+
     print("📦 Loading dataset...")
     dataset = load_dataset(dataset_name, split="train")
-    dataset = dataset.select(range(min(max_encode, len(dataset))))
+    limit = min(max_encode, len(dataset))
+    dataset = dataset.select(range(limit))
+    print(f"⚡ Using {len(dataset)} prompts for training")
 
-    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
+    # Split train/test
+    train_size = int(0.9 * len(dataset))
+    dataset_train = dataset.select(range(train_size))
+    dataset_test = dataset.select(range(train_size, len(dataset)))
+    print(f"🧪 Train/Test split: {len(dataset_train)} / {len(dataset_test)}")
 
-    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:
-            print(f"  → Encoded {i+1}/{len(dataset)}")
+    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
 
-    short_embeddings = torch.vstack(short_list)
-    long_embeddings = torch.vstack(long_list)
+    def encode_dataset(ds):
+        short_list, long_list = [], []
+        for i, item in enumerate(ds):
+            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(ds):
+                print(f"  → Encoded {i+1}/{len(ds)} prompts")
+                gc.collect()
+        return torch.vstack(short_list), torch.vstack(long_list)
 
-    input_dim = short_embeddings.shape[1]
-    output_dim = long_embeddings.shape[1]
-    model = GemmaTrainer(input_dim, 512, output_dim)
+    short_train, long_train = encode_dataset(dataset_train)
+    short_test, long_test = encode_dataset(dataset_test)
+    print(f"✅ Embeddings shapes: train {short_train.shape}/{long_train.shape}, test {short_test.shape}/{long_test.shape}")
 
-    criterion = nn.MSELoss()
+    input_dim = short_train.shape[1]
+    output_dim = long_train.shape[1]
+    model = GemmaTrainer(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim).to(device)
+
+    criterion = nn.CosineSimilarity(dim=1)
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
+    max_epochs = 50
+    batch_size = 32
+    n = short_train.shape[0]
 
     print("🚀 Training model...")
-    for epoch in range(10):
+    for epoch in range(max_epochs):
         model.train()
-        optimizer.zero_grad()
-        outputs = model(short_embeddings)
-        loss = criterion(outputs, long_embeddings)
-        loss.backward()
-        optimizer.step()
-        print(f"Epoch {epoch+1}/10 - Loss: {loss.item():.6f}")
+        epoch_loss = 0.0
+        perm = torch.randperm(n)
+        for start in range(0, n, batch_size):
+            idx = perm[start:start+batch_size]
+            inputs = short_train[idx].to(device)
+            targets = long_train[idx].to(device)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = 1 - criterion(outputs, targets).mean()
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item() * inputs.size(0)
+        epoch_loss /= n
 
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        save_path = os.path.join(tmp_dir, "model.flashpack")
-        model.save_flashpack(save_path)
-        repo = Repository(local_dir=tmp_dir, clone_from="rahul7star/FlashPack", use_auth_token=True)
-        repo.push_to_hub()
-        print("✅ Model pushed to Hugging Face Hub!")
+        model.eval()
+        with torch.no_grad():
+            test_outputs = model(short_test.to(device))
+            test_loss = 1 - criterion(test_outputs, long_test.to(device)).mean()
+
+        if epoch % 5 == 0 or epoch == max_epochs - 1:
+            print(f"Epoch {epoch+1}/{max_epochs} — Train Loss: {epoch_loss:.6f}, Test Loss: {test_loss:.6f}")
 
-    return model, dataset, embed_model, tokenizer, long_embeddings
+        if test_loss < 0.01:
+            print("🎯 Early stopping — loss threshold reached.")
+            break
+
+    print("✅ Training complete!")
+
+    if push_to_hub and test_loss < 0.05:
+        push_flashpack_model_to_hf(model, hf_repo)
+    else:
+        print("⚠️ Model not pushed — test loss not low enough.")
+
+    return model, dataset, embed_model, tokenizer, long_train
 
 
 # ============================================================
-# Try loading from HF first
+# 5️⃣ Load or train
 # ============================================================
-def get_flashpack_model(hf_repo="rahul7star/FlashPack", input_dim=768, output_dim=768):
+def get_flashpack_model(hf_repo="rahul7star/FlashPack", input_dim=1536, output_dim=1536):
     try:
         print(f"🔁 Loading FlashPack model from {hf_repo}...")
-        dummy_model = GemmaTrainer(input_dim, 512, output_dim)
+        dummy_model = GemmaTrainer(input_dim=input_dim, hidden_dim=1024, output_dim=output_dim)
         model = GemmaTrainer.from_flashpack(hf_repo, model=dummy_model)
+        model.eval()
         print("✅ Model loaded successfully.")
-        tokenizer, embed_model, _ = build_encoder("gpt2", max_length=32)
+        tokenizer, embed_model, _ = build_encoder("gpt2", max_length=128)
         return model, tokenizer, embed_model, None, None
     except Exception as e:
         print(f"⚠️ Load failed: {e}")
         print("⏬ Training new model...")
         return train_flashpack_model()
 
+# ============================================================
+# 6️⃣ Encode and Enhance
+# ============================================================
+@torch.no_grad()
+def encode_prompt(prompt, tokenizer, embed_model):
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
+                       padding="max_length", max_length=128).to(device)
+    last_hidden = embed_model(**inputs).last_hidden_state
+    mean_pool = last_hidden.mean(dim=1)
+    max_pool, _ = last_hidden.max(dim=1)
+    return torch.cat([mean_pool, max_pool], dim=1).cpu()
+
+def make_enhance_fn(model, long_embeddings, dataset, tokenizer, embed_model):
+    @torch.no_grad()
+    def fn(user_prompt, chat_history):
+        if model is None or dataset is None or long_embeddings is None:
+            return [{"role": "system", "content": "⚠️ Model not loaded. Please train or reload first."}]
+        chat_history = chat_history or []
+        short_emb = encode_prompt(user_prompt, tokenizer, embed_model)
+        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))
+        best_idx = int(sims.argmax().item())
+        enhanced_prompt = dataset[best_idx]["long_prompt"]
+        chat_history.append({"role": "user", "content": user_prompt})
+        chat_history.append({"role": "assistant", "content": enhanced_prompt})
+        return chat_history
+    return fn
 
 # ============================================================
-# Gradio UI
+# 7️⃣ Gradio UI with Train Button
 # ============================================================
-with gr.Blocks(title="FlashPack Trainer") as demo:
-    gr.Markdown("## ⚡ FlashPack Trainer")
+with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ✨ Prompt Enhancer (FlashPack mapper)")
 
-    status = gr.Textbox(label="Status", value="Ready", interactive=False)
-    train_btn = gr.Button("Train")
+    status = gr.Textbox(value="Model loading...", label="Status", interactive=False)
+
+    chatbot = gr.Chatbot(label="Enhanced Prompts", type="messages", height=400)
+    user_prompt = gr.Textbox(placeholder="Enter a short prompt...", label="Your Prompt", lines=3)
+    send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
+    train_btn = gr.Button("🧠 Train Model")
+    clear_btn = gr.Button("🧹 Clear Chat")
 
     model_state = gr.State(None)
     tokenizer_state = gr.State(None)
@@ -142,15 +235,10 @@ with gr.Blocks(title="FlashPack Trainer") as demo:
     long_embeddings_state = gr.State(None)
 
     def train_model():
+        status_text = "🔄 Training or loading model..."
         model, tokenizer, embed_model, dataset, long_embeddings = get_flashpack_model()
-        return (
-            model,
-            tokenizer,
-            embed_model,
-            dataset,
-            long_embeddings,
-            "✅ Model ready for use",
-        )
+        status_text = "✅ Model ready!"
+        return model, tokenizer, embed_model, dataset, long_embeddings, status_text
 
     train_btn.click(
         train_model,
@@ -158,4 +246,32 @@ with gr.Blocks(title="FlashPack Trainer") as demo:
         outputs=[model_state, tokenizer_state, embed_model_state, dataset_state, long_embeddings_state, status],
     )
 
-demo.launch()
+    def enhance(user_prompt, chat_history, model, tokenizer, embed_model, dataset, long_embeddings):
+        fn = make_enhance_fn(model, long_embeddings, dataset, tokenizer, embed_model)
+        return fn(user_prompt, chat_history)
+
+    send_btn.click(
+        enhance,
+        inputs=[user_prompt, chatbot, model_state, tokenizer_state, embed_model_state, dataset_state, long_embeddings_state],
+        outputs=chatbot,
+    )
+    user_prompt.submit(
+        enhance,
+        inputs=[user_prompt, chatbot, model_state, tokenizer_state, embed_model_state, dataset_state, long_embeddings_state],
+        outputs=chatbot,
+    )
+
+    clear_btn.click(lambda: [], None, chatbot)
+
+    # Auto-load model on startup
+    demo.load(
+        lambda: train_model(),
+        inputs=None,
+        outputs=[model_state, tokenizer_state, embed_model_state, dataset_state, long_embeddings_state, status],
+    )
+
+# ============================================================
+# 9️⃣ Launch
+# ============================================================
+if __name__ == "__main__":
+    demo.launch(show_error=True)