Update functions/modelling_function.py

functions/modelling_function.py

@@ -78,7 +78,7 @@ def category_reassign(row, reference_df, checked_category, threshold=70):
     else:
         return row['category_name']
     
-def train_model(df, stratify=True, model_type='bert', use_existing_model=False, model_name=None):
+def train_model(df, train_type, label_column, stratify=True, model_type='bert', use_existing_model=False, model_name=None):
     """
     This function trains the model using the configuration in config.yaml
 
@@ -98,7 +98,7 @@ def train_model(df, stratify=True, model_type='bert', use_existing_model=False,
     warnings.filterwarnings('ignore')
     
     test_size = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)['parameters']['training_args']['test_size']
-    train_df, test_df = train_test_split(df, test_size=test_size, stratify=df['category_name'])
+    train_df, test_df = train_test_split(df, test_size=test_size, stratify=df[label_column])
 
     # Optional model configuration  
     model_config = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)['parameters']['model_args']
@@ -112,7 +112,7 @@ def train_model(df, stratify=True, model_type='bert', use_existing_model=False,
 
     # Create a ClassificationModel
     model_detail = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)['parameters']['model_types']
-    class_names = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)['parameters']['class_names']
+    class_names = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)['parameters']['class_names'][train_type]
 
     if use_existing_model:
         model = ClassificationModel(model_type, model_name, num_labels=len(class_names), args=model_args, use_cuda=False)
@@ -125,7 +125,7 @@ def train_model(df, stratify=True, model_type='bert', use_existing_model=False,
     # Evaluate the model
     result, model_outputs, wrong_predictions = model.eval_model(test_df)
     preds = np.argmax(model_outputs, axis=1)
-    class_report =classification_report(test_df['category_name'], preds, target_names=class_names)
+    class_report =classification_report(test_df[label_column], preds, target_names=class_names)
 
     return model, preds, class_report, train_df, test_df, class_names