Update app_flash.py

app_flash.py

@@ -9,7 +9,6 @@ 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
 
 # ============================================================
@@ -23,17 +22,14 @@ print(f"🔧 Using device: {device} (CPU-only)")
 # 1️⃣ FlashPack model with better hidden layers
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self):
+    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
         super().__init__()
-        input_dim = 1536
-        hidden_dim = 1024
-        output_dim = 1536
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
         self.fc2 = nn.Linear(hidden_dim, hidden_dim)
         self.fc3 = nn.Linear(hidden_dim, output_dim)
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.fc1(x)
         x = self.relu(x)
         x = self.fc2(x)
@@ -157,53 +153,6 @@ def train_flashpack_model(
 # 5️⃣ Load or train model
 # ============================================================
 def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
-    input_dim = 1536  # must match the input_dim used during training
-    try:
-        print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
-
-        # 1️⃣ Try local model first
-        local_model_path = "model.flashpack"
-        if os.path.exists(local_model_path):
-            print("✅ Loading local model")
-        else:
-            # 2️⃣ Try Hugging Face
-            files = list_repo_files(hf_repo)
-            if "model.flashpack" in files:
-                print("✅ Downloading model from HF")
-                from huggingface_hub import hf_hub_download
-                local_model_path = hf_hub_download(repo_id=hf_repo, filename="model.flashpack")
-            else:
-                print("🚫 No pretrained model found")
-                return None, None, None, None
-
-        # 3️⃣ Load model with correct input_dim
-        model = GemmaTrainer(input_dim=input_dim).from_flashpack(local_model_path)
-        model.eval()
-
-        # 4️⃣ Build encoder
-        tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=128)
-
-        # 5️⃣ Enhancement function
-        @torch.no_grad()
-        def enhance_fn(prompt, chat):
-            chat = chat or []
-            short_emb = encode_fn(prompt).to(device)
-            mapped = model(short_emb).cpu()
-            long_prompt = f"🌟 Enhanced prompt: {prompt} (creatively expanded)"
-            chat.append({"role": "user", "content": prompt})
-            chat.append({"role": "assistant", "content": long_prompt})
-            return chat
-
-        return model, tokenizer, embed_model, enhance_fn
-
-    except Exception as e:
-        print(f"⚠️ Load failed: {e}")
-        print("⏬ Training a new FlashPack model locally...")
-        model, dataset, embed_model, tokenizer, long_embeddings = train_flashpack_model()
-        push_flashpack_model_to_hf(model, hf_repo, log_fn=print)
-        return model, tokenizer, embed_model, None
-
-def get_flashpack_model1(hf_repo="rahul7star/FlashPack"):
     try:
         print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
         model = GemmaTrainer.from_flashpack(hf_repo)
@@ -280,4 +229,4 @@ with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft
 # 9️⃣ Launch
 # ============================================================
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(show_error=True)