Repository: davanstrien/flyswot (language: python)

Issue:
Add support for multi-task flyswot models
closes #154 
- adds support for multi-task/head models
	- adds support for printing reports containing multiple label groups
	- adds support for writing multiple labels groups to csv
- improved test coverage for inference methods


Discussion / hints:

Fix patch:
diff --git a/src/flyswot/inference.py b/src/flyswot/inference.py
index 36835f80..4eb46061 100644
--- a/src/flyswot/inference.py
+++ b/src/flyswot/inference.py
@@ -161,22 +161,13 @@ def predict_directory(
     typer.echo(f"CSV report stored in {csv_fname}")
     delta = timedelta(seconds=time.perf_counter() - start_time)
     typer.echo(f"Time taken to run:  {str(delta)}")
-    if len(onnxinference.vocab) > 1:
-        labels_to_print = labels_from_csv(csv_fname)
-        tables = [
-            print_table(labels, f"Prediction summary {i+1}", print=False)
-            for i, labels in enumerate(labels_to_print)
-        ]
-        columns = Columns(tables)
-        print(columns)
-    else:
-        labels_to_print = labels_from_csv(csv_fname)
-        tables = [
-            print_table(labels, f"Prediction summary {i+1}", print=False)
-            for i, labels in enumerate(labels_to_print)
-        ]
-        columns = Columns(tables)
-        print(columns)
+    labels_to_print = labels_from_csv(csv_fname)
+    tables = [
+        print_table(labels, f"Prediction summary {i+1}", print=False)
+        for i, labels in enumerate(labels_to_print)
+    ]
+    columns = Columns(tables)
+    print(columns)
     # print_table(list(itertoolz.concat(all_preds)))
 
 


Test patch:
diff --git a/tests/test_inference.py b/tests/test_inference.py
index 9dd74787..3cf770fd 100644
--- a/tests/test_inference.py
+++ b/tests/test_inference.py
@@ -134,9 +134,9 @@ def test_predict_directory(datafiles, tmp_path) -> None:
         for row in reader:
             for (k, v) in row.items():
                 columns[k].append(v)
-        assert any(["prediction" in k for k in columns])
-        labels = [columns[k] for k in columns.keys() if "prediction" in k]
-        confidences = [columns[k] for k in columns.keys() if "confidence" in k]
+        assert any("prediction" in k for k in columns)
+        labels = [columns[k] for k in columns if "prediction" in k]
+        confidences = [columns[k] for k in columns if "confidence" in k]
         # check all labels are strings
         assert all(map(lambda x: isinstance(x, str), (itertoolz.concat(labels))))
         # check all confidences can be cast to float