Update app_flash1.py

app_flash1.py

@@ -21,9 +21,17 @@ print(f"🔧 Using device: {device} (CPU-only mode)")
 # ============================================================
 # 1️⃣ Model Definition
 # ============================================================
+# ============================================================
+# 1️⃣ Fixed Model Definition (FlashPack-compatible)
+# ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
+    def __init__(self):
         super().__init__()
+        # Move input/output dimensions inside
+        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)
@@ -38,6 +46,7 @@ class GemmaTrainer(nn.Module, FlashPackMixin):
         return x
 
 
+
 # ============================================================
 # 2️⃣ Encoder Setup
 # ============================================================
@@ -131,10 +140,9 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     # 1️⃣ Try local first
     if os.path.exists(local_path):
         print("✅ Found local model.flashpack — loading it directly.")
-        model = GemmaTrainer().from_flashpack(local_path)
+        model = GemmaTrainer().from_flashpack(local_path)  # no args needed
         model.eval()
         tokenizer, embed_model, _ = build_encoder("gpt2")
-        print("✅ FlashPack model loaded successfully (local).")
         return model, tokenizer, embed_model, None, None
 
     # 2️⃣ Otherwise, check HF repo
@@ -143,19 +151,13 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
         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")
-
-            # Correct loading: no extra args passed
-            model = GemmaTrainer().from_flashpack(local_path)
+            model = GemmaTrainer().from_flashpack(local_path)  # no args
             model.eval()
-
             tokenizer, embed_model, _ = build_encoder("gpt2")
-            print("✅ FlashPack model loaded successfully (from repo).")
             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"⚠️ Error checking repo: {e}")
         print("⏬ Training new model instead.")
@@ -164,6 +166,7 @@ def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
 
 
 
+
 # ============================================================
 # 6️⃣ Encode & Enhance Functions
 # ============================================================