typing fix

CogStack · Mar 5, 2024 · 404a246 · 404a246
1 parent 493c0ff
commit 404a246
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/medcat/evaluate_mct_export/mct_analysis.py b/medcat/evaluate_mct_export/mct_analysis.py
@@ -350,7 +350,14 @@ def full_annotation_df(self) -> pd.DataFrame:
                 else:
                     pred_meta_values.append(_meta_values.get(meta_results['predictions'][counter], np.nan))
                     counter += 1
-            meta_df.insert(int(meta_df.columns.get_loc(meta_model)) + 1, f'predict_{meta_model}', pred_meta_values) # TODO fix this line
+
+            loc = meta_df.columns.get_loc(meta_model)
+            if isinstance(loc, int):
+                meta_df.insert(loc + 1, f'predict_{meta_model}', pred_meta_values)
+            else:
+                print(f"Warning: Unexpected column location type: {type(loc)}")
+            meta_df.insert(1, f'predict_{meta_model}', pred_meta_values)
+            #meta_df.insert(int(meta_df.columns.get_loc(meta_model)) + 1, f'predict_{meta_model}', pred_meta_values) # TODO fix this line
 
         return meta_df
 

diff --git a/tests/medcat/evaluate_mct_export/offline_test_mct_analysis.py b/tests/medcat/evaluate_mct_export/offline_test_mct_analysis.py
@@ -29,7 +29,6 @@ def setUpClass(cls) -> None:
         cls.export = MedcatTrainer_export([MCT_EXPORT_JSON_PATH, ], MODEL_PACK_PATH)
 
     # these would need a CAT instance
-    ''' TODO: return this once meta_annotations are fixed
     def test_can_full_annotation_df(self):
         full_ann_df = self.export.full_annotation_df()
         self.assertNonEmptyDataframe(full_ann_df)
@@ -39,7 +38,6 @@ def test_can_meta_anns_concept_summary(self):
         # this will be empty since I don't think I have anything
         # of note regarding meta annotations
         self.assertIsNotNone(meta_anns_summary_df)
-    '''
 
     def test_generate_report(self):
         self.export.generate_report(path=self.report_path)