Update app.py

app.py

@@ -48,53 +48,122 @@ keras_model = None
 kmer_to_index = None
 analyzer = None
 
-# --- Load Models ---
-def load_models():
-    global boundary_model, keras_model, kmer_to_index
-    hf_token = os.getenv("HF_TOKEN", None)
+ --- Load Models ---
+boundary_model = None
+keras_model = None
+kmer_to_index = None
+
+# Try to load boundary model from Hugging Face Hub
+try:
+    boundary_path = hf_hub_download(
+        repo_id=model_repo, 
+        filename="best_boundary_aware_model.pth",
+        token=hf_token
+    )
+    if os.path.exists(boundary_path):
+        boundary_model = GenePredictor(boundary_path)
+        logging.info("Boundary model loaded successfully from Hugging Face Hub.")
+    else:
+        logging.warning(f"Boundary model file not found after download")
+except Exception as e:
+    logging.error(f"Failed to load boundary model from HF Hub: {e}")
+
+# Try to load Keras model from Hugging Face Hub
+try:
+    keras_path = hf_hub_download(
+        repo_id=model_repo, 
+        filename="best_model.keras",
+        token=hf_token
+    )
+    kmer_path = hf_hub_download(
+        repo_id=model_repo, 
+        filename="kmer_to_index.pkl",
+        token=hf_token
+    )
     
-    # Load boundary model
-    if GenePredictor:
-        try:
-            boundary_path = hf_hub_download(
-                repo_id=MODEL_REPO,
-                filename="best_boundary_aware_model.pth",
-                token=hf_token,
-                cache_dir="/data/models"
-            )
-            boundary_model = GenePredictor(boundary_path)
-            logging.info("Boundary model loaded successfully.")
-        except Exception as e:
-            logging.warning(f"Failed to load boundary model: {e}")
-            boundary_model = None
+    if os.path.exists(keras_path) and os.path.exists(kmer_path):
+        keras_model = load_model(keras_path)
+        with open(kmer_path, "rb") as f:
+            kmer_to_index = pickle.load(f)
+        logging.info("Keras model and k-mer index loaded successfully from Hugging Face Hub.")
     else:
-        logging.warning("GenePredictor not available.")
+        logging.warning(f"Keras model or kmer files not found after download")
+except Exception as e:
+    logging.error(f"Failed to load Keras model from HF Hub: {e}")
+
+# --- Load Verification Models from models directory ---
+verification_models = {}
+
+def load_verification_models():
+    """Load all verification models from the models directory"""
+    global verification_models
+    models_dir = "models"
+    
+    if not os.path.exists(models_dir):
+        logging.warning(f"Models directory not found: {models_dir}")
+        return
+    
+    # Load different types of verification models
+    model_files = {
+        "boundary_model": "best_boundary_aware_model.pth",
+        "keras_model": "best_model.keras", 
+        "kmer_index": "kmer_to_index.pkl",
+        "additional_model_1": "verification_model_1.pth",  # Add your model names here
+        "additional_model_2": "verification_model_2.keras",
+        # Add more models as needed
+    }
     
-    # Load Keras model
-    if load_model:
+    for model_name, filename in model_files.items():
+        model_path = os.path.join(models_dir, filename)
+        
         try:
-            keras_path = hf_hub_download(
-                repo_id=MODEL_REPO,
-                filename="best_model.keras",
-                token=hf_token,
-                cache_dir="/data/models"
-            )
-            kmer_path = hf_hub_download(
-                repo_id=MODEL_REPO,
-                filename="kmer_to_index.pkl",
-                token=hf_token,
-                cache_dir="/data/models"
-            )
-            keras_model = load_model(keras_path)
-            with open(kmer_path, "rb") as f:
-                kmer_to_index = pickle.load(f)
-            logging.info("Keras model and k-mer index loaded successfully.")
+            if os.path.exists(model_path):
+                if filename.endswith('.pth'):
+                    # PyTorch model
+                    if model_name == "boundary_model":
+                        verification_models[model_name] = GenePredictor(model_path)
+                    else:
+                        verification_models[model_name] = torch.load(model_path, map_location='cpu')
+                        
+                elif filename.endswith('.keras'):
+                    # Keras model
+                    verification_models[model_name] = load_model(model_path)
+                    
+                elif filename.endswith('.pkl'):
+                    # Pickle file
+                    with open(model_path, 'rb') as f:
+                        verification_models[model_name] = pickle.load(f)
+                
+                logging.info(f"Loaded verification model: {model_name}")
+                
         except Exception as e:
-            logging.warning(f"Failed to load Keras model or k-mer index: {e}")
-            keras_model = None
-            kmer_to_index = None
+            logging.error(f"Failed to load {model_name} from {model_path}: {e}")
+
+# Load verification models at startup
+load_verification_models()
+
+# --- Initialize Tree Analyzer ---
+analyzer = None
+try:
+    analyzer = ml_simplified_tree.PhylogeneticTreeAnalyzer()
+    if os.path.exists(csv_path):
+        if analyzer.load_data(csv_path):
+            logging.info("Tree analyzer initialized successfully")
+            # Try to train AI model (optional)
+            try:
+                if not analyzer.train_ai_model():
+                    logging.warning("AI model training failed; proceeding with basic analysis.")
+            except Exception as e:
+                logging.warning(f"AI model training failed: {e}")
+        else:
+            logging.error("Failed to load CSV data for tree analyzer")
+            analyzer = None
     else:
-        logging.warning("Keras/TensorFlow not available.")
+        logging.error(f"CSV file not found: {csv_path}")
+        analyzer = None
+except Exception as e:
+    logging.error(f"Failed to initialize tree analyzer: {e}")
+    analyzer = None
 
 # --- Initialize Tree Analyzer ---
 def init_tree_analyzer():