plot_mean_test_f1: add --restricted plot

Author: Aethor

plot_mean_test_f1.py

@@ -7,6 +7,8 @@
 parser = argparse.ArgumentParser()
 parser.add_argument("-o", "--output", type=str, default=None)
 parser.add_argument("-r", "--oracle", action="store_true")
+parser.add_argument("-e", "--restricted", action="store_true")
+parser.add_argument("--no-baseline", action="store_true")
 args = parser.parse_args()
 
 # from https://matplotlib.org/stable/gallery/lines_bars_and_markers/linestyles.html
@@ -29,7 +31,19 @@
 ]
 
 
-runs = ["random", "bm25", "samenoun", "left", "right", "neighbors"]
+# runs is of form {run_dir_name => name}
+if args.oracle:
+    runs = {
+        f"oracle_{run}": run
+        for run in ["random", "bm25", "samenoun", "before", "after", "surrounding"]
+    }
+elif args.restricted:
+    runs = {"oracle_bm25": "bm25", "bm25_restricted": "restricted bm25"}
+else:
+    runs = {
+        run: run
+        for run in ["random", "bm25", "samenoun", "before", "after", "surrounding"]
+    }
 
 with open(f"./runs/bare/metrics.json") as f:
     bare_metrics = json.load(f)
@@ -42,9 +56,7 @@
 
 fig.set_size_inches(16, 12)
 
-for run_i, run in enumerate(runs):
-    if args.oracle:
-        run = f"oracle_{run}"
+for run_i, (run, run_name) in enumerate(runs.items()):
     with open(f"./runs/short/{run}/metrics.json") as f:
         metrics = json.load(f)
     ax.plot(
@@ -54,19 +66,20 @@
         linewidth=4,
     )
 
-# bare baseline
-ax.plot(
-    [1, 6],
-    [bare_metrics["mean_test_f1"]["values"][0]] * 2,
-    linestyle=linestyle_tuple[len(runs)][1],
-    linewidth=4,
-)
+# no retrieval baseline
+if not args.no_baseline:
+    ax.plot(
+        [1, 6],
+        [bare_metrics["mean_test_f1"]["values"][0]] * 2,
+        linestyle=linestyle_tuple[len(runs)][1],
+        linewidth=4,
+    )
 
 ax.grid()
 ax.set_ylabel("F1", fontsize=40)
 ax.set_xlabel("Number of retrieved sentences", fontsize=40)
 ax.legend(
-    runs + ["no retrieval"],
+    list(runs.values()) + ["no retrieval"] if not args.no_baseline else [],
     loc="lower center",
     bbox_to_anchor=(0.5, 1),
     fontsize=40,