Fix import

.gitignore

@@ -1,3 +1,4 @@
 *__pycache__/
 example_event.html
 web/
+mlruns/

src/predict/main.py

@@ -16,9 +16,9 @@ from .models import (
 MODELS_TO_RUN = [
     EloBaselineModel(),
     LogisticRegressionModel(),
-    XGBoostModel(),
-    SVCModel(),
-    RandomForestModel(),
+    # XGBoostModel(),
+    # SVCModel(),
+    # RandomForestModel(),
     BernoulliNBModel(),
     LGBMModel(),
 ]
@@ -58,7 +58,7 @@ def main():
     parser.add_argument(
         '--kfold',
         action='store_true',
-        help='Run 3-fold CV instead of standard split.'
+        help='Run 10-fold CV instead of standard split.'
     )
     args = parser.parse_args()
 

src/predict/models.py

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 import sys
 import os
 import pandas as pd
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, List
 from sklearn.linear_model import LogisticRegression
 from sklearn.svm import SVC
 from sklearn.naive_bayes import BernoulliNB
@@ -88,7 +88,7 @@ class BaseMLModel(BaseModel):
         self.fighters_df = None
         self.fighter_histories = {}
 
-    def train(self, train_fights: list[dict[str, any]]) -> None:
+    def train(self, train_fights: List[Dict[str, Any]]) -> None:
         """
         Trains the machine learning model. This involves loading fighter data,
         pre-calculating histories, and fitting the model on the preprocessed data.

src/predict/pipeline.py

@@ -291,27 +291,17 @@ class PredictionPipeline:
                     # Train and evaluate
                     model.train(train_set)
                     correct = 0
-                    total_fights = 0
                     for fight in test_set:
-                        if fight['winner'] not in ["Draw", "NC", ""]:
-                            prediction = model.predict(fight)
-                            if prediction.get('winner') == fight['winner']:
-                                correct += 1
-                            total_fights += 1
+                        prediction = model.predict(fight)
+                        if prediction.get('winner') == fight['winner']:
+                            correct += 1
 
-                    acc = correct / total_fights if total_fights > 0 else 0.0
+                    acc = correct / len(test_set) if test_set else 0.0
                     fold_results[model_name] = acc
 
-                    # Log metrics and register model to appear in MLflow Models tab
+                    # Log metrics and model artifact
                     mlflow.log_metric(f"accuracy_{model_name}", acc)
-                    mlflow.log_metric(f"total_fights_{model_name}", total_fights)
-                    
-                    # Register the model with MLflow to appear in Models tab
-                    mlflow.sklearn.log_model(
-                        model, 
-                        f"model_{model_name}",
-                        registered_model_name=f"{model_name}_UFC_Model"
-                    )
+                    mlflow.sklearn.log_model(model, f"model_{model_name}")
 
                 all_fold_metrics.append(fold_results)