model training pipeline modifs

MartinKalema · Jun 11, 2024 · f7738ef · f7738ef
1 parent 0751a06
commit f7738ef
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/swahiliNewsClassifier/components/model_training_and_evaluation.py b/src/swahiliNewsClassifier/components/model_training_and_evaluation.py
@@ -27,7 +27,7 @@ def __init__(self, model_training_and_evaluation_configurations: ModelTrainingAn
         """
         self.model_training_and_evaluation_configurations = model_training_and_evaluation_configurations
         self.bucket_name = "swahili-news-classifier"
-        self.model_path = f"models/text_classifier_learner.pth"
+        self.model_path = f"models/text_classifier_learner.pkl"
         self.s3 = boto3.client('s3', aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'), aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'), region_name=os.getenv('REGION_NAME'))
 
     def upload_to_s3(self) -> None:
@@ -158,7 +158,7 @@ def train_text_classifier(self, dls) -> None:
         learn.fit_one_cycle(self.model_training_and_evaluation_configurations.epochs_5, slice(1e-3/(2.6**4), self.model_training_and_evaluation_configurations.learning_rate_5))
         classifier_metrics = learn.validate()
         self.log_to_mlflow(classifier_metrics)
-        learn.save_encoder(f'text_classifier_learner')
+        learn.export('models/text_classifier_learner.pkl')
 
 
 

diff --git a/src/swahiliNewsClassifier/pipeline/stage_02_model_training_and_evaluation.py b/src/swahiliNewsClassifier/pipeline/stage_02_model_training_and_evaluation.py
@@ -19,7 +19,7 @@ def main(self):
         try:
             model_training_and_evaluation_config = self.config.get_model_training_and_evaluation_config()
             model_training_and_evaluation = ModelTrainingAndEvaluation(
-                model_training_and_evaluation_config=model_training_and_evaluation_config)
+                model_training_and_evaluation_configurations=model_training_and_evaluation_config)
             model_training_and_evaluation.run_pipeline()
         except Exception as e:
             log.exception(f"An error occurred during {STAGE_NAME}: {e}")