tests and autodetection for AI pipeline mode

Author: dsmutin

.gitignore

@@ -18,7 +18,7 @@ data/databases/open/test_ML_database_tokenized.csv
 genome_parse_results.csv
 data/genomes/
 data/articles2/
-data/test/
+data/test/*
 benchmark*
 OligoMiner
 

pipeline.py

@@ -42,6 +42,7 @@
 from PROBESt.modeling import run_modeling
 from PROBESt.dedegeneration import run_dedegeneration
 from PROBESt.prepare_blast import prepare_bases_if_needed
+from PROBESt.AI import apply_ai_filtration_to_blast_file
 
 # Functions
 
@@ -163,6 +164,23 @@ def merge_iter(iter: int):
 
     print("Negative hits counted")
 
+    # 2.5. like domain supersecondary organization ----
+    if getattr(args, 'AI', True):  # Default to True if not set
+        model_path = 'data/AIL.pt'
+        positive_hits_path = out_dir(iter) + "positive_hits.tsv"
+        input_size = 14
+        
+        # Apply AI filtration (function handles all error checking and logging)
+        # model_architecture=None will auto-detect from saved model
+        apply_ai_filtration_to_blast_file(
+            positive_hits_path=positive_hits_path,
+            model_path=model_path,
+            fasta_file=prev_merged_fa,
+            input_size=input_size,
+            model_architecture=None,  # Auto-detect from saved model
+            threshold_method="median"
+        )
+
     # 3. multimapping detection and filtering ----
     probe_check_iter = probe_check + \
         " -o " + out_dir(iter) + "clear_hits.tsv" + \

scripts/generator/ML_filtration.py

@@ -358,13 +358,29 @@ def main():
 
         # Save the model
         os.makedirs('data', exist_ok=True)
+        # Determine model architecture name from best_model_name
+        model_architecture = best_model_name
+        # Map common names to architecture names
+        if "GAIL_Wide_Extra" in best_model_name or "GAILWideExtra" in best_model_name:
+            model_architecture = "GAILWideExtra"
+        elif "GAIL_Wide" in best_model_name:
+            model_architecture = "GAILWide"
+        elif "GAIL_Deep" in best_model_name:
+            model_architecture = "GAILDeep"
+        elif "GAIL_Narrow" in best_model_name:
+            model_architecture = "GAILNarrow"
+        elif "GAIL" in best_model_name:
+            model_architecture = "GAILDiscriminator"
+        
         torch.save({
             'model_state_dict': best_model.model.state_dict(),
             'scaler': best_model.scaler,
             'learning_rate': best_model.learning_rate,
-            'weight_pos': best_model.criterion.pos_weight.item() if hasattr(best_model.criterion, 'pos_weight') else None
+            'weight_pos': best_model.criterion.pos_weight.item() if hasattr(best_model.criterion, 'pos_weight') else None,
+            'model_architecture': model_architecture,
+            'model_name': best_model_name
         }, 'data/GAIL.pt')
-        print(f"Model saved to data/GAIL.pt")
+        print(f"Model saved to data/GAIL.pt (architecture: {model_architecture})")
 
         # Re-validate after extended training
         final_val_metrics = validate_filtration_AI(