Update app_flash1.py

app_flash1.py

@@ -1,5 +1,5 @@
-import gc
 import os
+import gc
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -8,18 +8,18 @@ import gradio as gr
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModel
 from flashpack import FlashPackMixin
-from huggingface_hub import Repository
+from huggingface_hub import Repository, list_repo_files, hf_hub_download
 from typing import Tuple
 
 # ============================================================
-# 🖥 Device setup (CPU-only)
+# 🖥 Device Setup
 # ============================================================
 device = torch.device("cpu")
 torch.set_num_threads(4)
 print(f"🔧 Using device: {device} (CPU-only mode)")
 
 # ============================================================
-# 1️⃣ FlashPack model
+# 1️⃣ Model Definition
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
     def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
@@ -29,7 +29,7 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         self.fc2 = nn.Linear(hidden_dim, hidden_dim)
         self.fc3 = nn.Linear(hidden_dim, output_dim)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor):
         x = self.fc1(x)
         x = self.relu(x)
         x = self.fc2(x)
@@ -37,241 +37,174 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         x = self.fc3(x)
         return x
 
+
 # ============================================================
-# 2️⃣ Encoder (mean + max pooling)
+# 2️⃣ Encoder Setup
 # ============================================================
-def build_encoder(model_name="gpt2", max_length: int = 128):
-    print(f"📦 Loading tokenizer and model for {model_name}...")
+def build_encoder(model_name="gpt2", max_length=128):
+    print(f"📦 Loading encoder: {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).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).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
+        hidden = embed_model(**inputs).last_hidden_state
+        mean_pool = hidden.mean(dim=1)
+        max_pool, _ = hidden.max(dim=1)
+        return torch.cat([mean_pool, max_pool], dim=1).cpu()
 
     return tokenizer, embed_model, encode
 
+
 # ============================================================
-# 3️⃣ Push FlashPack model to HF
+# 3️⃣ Push to Hugging Face
 # ============================================================
-def push_flashpack_model_to_hf(model, hf_repo: str):
-    logs = []
+def push_flashpack_model_to_hf(model, hf_repo):
     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.")
+        model.save_flashpack(os.path.join(tmp_dir, "model.flashpack"))
+        with open(os.path.join(tmp_dir, "README.md"), "w") as f:
+            f.write("# FlashPack Model\nTrained locally and pushed to HF.")
         repo.push_to_hub()
-        logs.append(f"✅ Model pushed to HF: {hf_repo}")
-    return logs
+        print(f"✅ Model pushed to {hf_repo}")
+
 
 # ============================================================
-# 4️⃣ Train FlashPack model
+# 4️⃣ Training Logic
 # ============================================================
-def train_flashpack_model(
-    dataset_name: str = "rahul7star/prompt-enhancer-dataset",
-    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]:
-
+def train_flashpack_model(dataset_name="rahul7star/prompt-enhancer-dataset",
+                          hf_repo="rahul7star/FlashPack",
+                          max_encode=1000):
     print("📦 Loading dataset...")
-    dataset = load_dataset(dataset_name, split="train")
-    limit = min(max_encode, len(dataset))
-    dataset = dataset.select(range(limit))
-    print(f"⚡ Using {len(dataset)} prompts for training")
-
-    # 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)}")
+    dataset = load_dataset(dataset_name, split="train").select(range(max_encode))
+    print(f"✅ Loaded {len(dataset)} samples")
 
-    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
+    tokenizer, embed_model, encode_fn = build_encoder("gpt2")
 
     def encode_dataset(ds):
-        short_list, long_list = [], []
+        s_list, l_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")
+            s_list.append(encode_fn(item["short_prompt"]))
+            l_list.append(encode_fn(item["long_prompt"]))
+            if (i + 1) % 50 == 0:
+                print(f"  → Encoded {i + 1}/{len(ds)}")
                 gc.collect()
-        return torch.vstack(short_list), torch.vstack(long_list)
-
-    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}")
+        return torch.vstack(s_list), torch.vstack(l_list)
 
-    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)
+    short_emb, long_emb = encode_dataset(dataset)
+    input_dim, output_dim = short_emb.shape[1], long_emb.shape[1]
+    model = GemmaTrainer(input_dim, 1024, output_dim)
 
-    criterion = nn.CosineSimilarity(dim=1)
     optimizer = optim.Adam(model.parameters(), lr=1e-3)
-    max_epochs = 50
-    batch_size = 32
-    n = short_train.shape[0]
+    loss_fn = nn.CosineSimilarity(dim=1)
 
     print("🚀 Training model...")
-    for epoch in range(max_epochs):
+    for epoch in range(20):
         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_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
-
-        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}")
-
-        if test_loss < 0.01:
-            print("🎯 Early stopping — loss threshold reached.")
+        optimizer.zero_grad()
+        preds = model(short_emb)
+        loss = 1 - loss_fn(preds, long_emb).mean()
+        loss.backward()
+        optimizer.step()
+        print(f"Epoch {epoch+1}/20 | Loss: {loss.item():.5f}")
+        if loss.item() < 0.01:
+            print("🎯 Early stopping.")
             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
+    push_flashpack_model_to_hf(model, hf_repo)
+    return model, tokenizer, embed_model, dataset, long_emb
 
 
 # ============================================================
-# 5️⃣ Load or train
+# 5️⃣ Load or Train
 # ============================================================
-def get_flashpack_model(hf_repo="rahul7star/FlashPack", input_dim=1536, output_dim=1536):
+def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
+    print(f"🔍 Checking for model in repo: {hf_repo}")
     try:
-        print(f"🔁 Loading FlashPack model from {hf_repo}...")
-        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=128)
-        return model, tokenizer, embed_model, None, None
+        files = list_repo_files(hf_repo)
+        if "model.flashpack" in files:
+            print("✅ Found model.flashpack in repo — downloading and loading it.")
+            local_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
+            dummy = GemmaTrainer(1536, 1024, 1536)
+            model = GemmaTrainer.from_flashpack(local_path, model=dummy)
+            model.eval()
+            tokenizer, embed_model, _ = build_encoder("gpt2")
+            return model, tokenizer, embed_model, None, None
+        else:
+            print("🚫 model.flashpack not found — starting training.")
+            return train_flashpack_model(hf_repo=hf_repo)
     except Exception as e:
-        print(f"⚠️ Load failed: {e}")
-        print("⏬ Training new model...")
-        return train_flashpack_model()
+        print(f"⚠️ Error checking repo: {e}")
+        print("⏬ Training new model instead.")
+        return train_flashpack_model(hf_repo=hf_repo)
+
 
 # ============================================================
-# 6️⃣ Encode and Enhance
+# 6️⃣ Encode & Enhance Functions
 # ============================================================
 @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)
+    hidden = embed_model(**inputs).last_hidden_state
+    mean_pool = hidden.mean(dim=1)
+    max_pool, _ = 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):
+
+def make_enhance_fn(model, tokenizer, embed_model, long_emb, dataset):
     @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)
+    def fn(prompt, chat):
+        chat = chat or []
+        short_emb = encode_prompt(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
+        sims = (long_emb @ mapped.t()).squeeze(1)
+        best = int(sims.argmax())
+        enhanced = dataset[best]["long_prompt"]
+        chat.append({"role": "user", "content": prompt})
+        chat.append({"role": "assistant", "content": enhanced})
+        return chat
     return fn
 
+
 # ============================================================
-# 7️⃣ Gradio UI with Train Button
+# 7️⃣ Gradio UI
 # ============================================================
-with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ✨ Prompt Enhancer (FlashPack mapper)")
+with gr.Blocks(title="✨ FlashPack Prompt Enhancer") as demo:
+    gr.Markdown("## 🧠 FlashPack Prompt Enhancer (CPU)\nShort → Long prompt expander")
 
-    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)
+    chatbot = gr.Chatbot(height=400)
+    user_input = gr.Textbox(label="Your prompt")
     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)
-    embed_model_state = gr.State(None)
-    dataset_state = gr.State(None)
-    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()
-        status_text = "✅ Model ready!"
-        return model, tokenizer, embed_model, dataset, long_embeddings, status_text
-
-    train_btn.click(
-        train_model,
-        inputs=[],
-        outputs=[model_state, tokenizer_state, embed_model_state, dataset_state, long_embeddings_state, status],
-    )
-
-    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 = gr.Button("🧹 Clear")
+    train_btn = gr.Button("🧩 Train Model", variant="secondary")
 
-    clear_btn.click(lambda: [], None, chatbot)
+    status = gr.Markdown("Status: Ready")
 
-    # 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],
-    )
+    # Load model initially
+    model, tokenizer, embed_model, dataset, long_emb = get_flashpack_model()
+    enhance_fn = make_enhance_fn(model, tokenizer, embed_model, long_emb, dataset) if dataset else None
+
+    def enhance(prompt, chat):
+        if not enhance_fn:
+            return chat + [{"role": "assistant", "content": "⚠️ Model not ready. Please train first."}]
+        return enhance_fn(prompt, chat)
+
+    def retrain():
+        global model, tokenizer, embed_model, dataset, long_emb, enhance_fn
+        model, tokenizer, embed_model, dataset, long_emb = train_flashpack_model()
+        enhance_fn = make_enhance_fn(model, tokenizer, embed_model, long_emb, dataset)
+        return "✅ Model retrained and pushed to HF!"
+
+    send_btn.click(enhance, [user_input, chatbot], chatbot)
+    user_input.submit(enhance, [user_input, chatbot], chatbot)
+    clear_btn.click(lambda: [], None, chatbot)
+    train_btn.click(retrain, None, status)
 
-# ============================================================
-# 9️⃣ Launch
-# ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)