From 9643dff294c66e6f60b630ccd6341f8b80e62f6f Mon Sep 17 00:00:00 2001 From: Alexey Taymanov Date: Thu, 16 Nov 2023 01:42:06 -0500 Subject: [PATCH] black --- .../test_pick_best_text_embedder.py | 67 +++++++++++-------- 1 file changed, 39 insertions(+), 28 deletions(-) diff --git a/tests/unit_tests/test_pick_best_text_embedder.py b/tests/unit_tests/test_pick_best_text_embedder.py index 02fc114..ea34db8 100644 --- a/tests/unit_tests/test_pick_best_text_embedder.py +++ b/tests/unit_tests/test_pick_best_text_embedder.py @@ -11,7 +11,9 @@ def test_pickbest_textembedder_missing_context_not_throws() -> None: featurizer = pick_best_chain.PickBestFeaturizer( auto_embed=False, model=MockEncoder() ) - event = pick_best_chain.PickBestEvent(inputs={"action": ToSelectFrom(["0", "1", "2"])}) + event = pick_best_chain.PickBestEvent( + inputs={"action": ToSelectFrom(["0", "1", "2"])} + ) featurizer.featurize(event) @@ -59,7 +61,7 @@ def test_pickbest_textembedder_w_label_no_score_no_emb() -> None: selected = pick_best_chain.PickBestSelected(index=0, probability=1.0) event = pick_best_chain.PickBestEvent( inputs={"context": BasedOn("context"), "action": ToSelectFrom(["0", "1", "2"])}, - selected=selected + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -105,8 +107,9 @@ def test_pickbest_textembedder_w_full_label_w_emb() -> None: event = pick_best_chain.PickBestEvent( inputs={ "context": rl_chain.Embed(BasedOn("ctx")), - "action": rl_chain.Embed(ToSelectFrom(["0", "1", "2"])) - }, selected=selected + "action": rl_chain.Embed(ToSelectFrom(["0", "1", "2"])), + }, + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -132,8 +135,9 @@ def test_pickbest_textembedder_w_full_label_w_embed_and_keep() -> None: event = pick_best_chain.PickBestEvent( inputs={ "context": rl_chain.EmbedAndKeep(BasedOn("ctx")), - "action": rl_chain.EmbedAndKeep(ToSelectFrom(["0", "1", "2"])) - }, selected=selected + "action": rl_chain.EmbedAndKeep(ToSelectFrom(["0", "1", "2"])), + }, + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -155,8 +159,8 @@ def test_pickbest_textembedder_more_namespaces_no_label_no_emb() -> None: inputs={ "action": ToSelectFrom([{"a": "0", "b": "0"}, "1", "2"]), "context1": BasedOn("context1"), - "context2": BasedOn("context2") - } + "context2": BasedOn("context2"), + } ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -179,9 +183,9 @@ def test_pickbest_textembedder_more_namespaces_w_label_no_emb() -> None: inputs={ "action": ToSelectFrom([{"a": "0", "b": "0"}, "1", "2"]), "context1": BasedOn("context1"), - "context2": BasedOn("context2") - }, - selected=selected + "context2": BasedOn("context2"), + }, + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -204,9 +208,9 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_no_emb() -> None: inputs={ "action": ToSelectFrom([{"a": "0", "b": "0"}, "1", "2"]), "context1": BasedOn("context1"), - "context2": BasedOn("context2") - }, - selected=selected + "context2": BasedOn("context2"), + }, + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -236,9 +240,9 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_emb() -> None inputs={ "context1": BasedOn(rl_chain.Embed(ctx_str_1)), "context2": BasedOn(rl_chain.Embed(ctx_str_2)), - "action": ToSelectFrom(rl_chain.Embed([{"a": "0", "b": "0"}, "1", "2"])) + "action": ToSelectFrom(rl_chain.Embed([{"a": "0", "b": "0"}, "1", "2"])), }, - selected=selected + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -270,9 +274,11 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_embed_and_kee inputs={ "context1": BasedOn(rl_chain.EmbedAndKeep(ctx_str_1)), "context2": BasedOn(rl_chain.EmbedAndKeep(ctx_str_2)), - "action": ToSelectFrom(rl_chain.EmbedAndKeep([{"a": "0", "b": "0"}, "1", "2"])) + "action": ToSelectFrom( + rl_chain.EmbedAndKeep([{"a": "0", "b": "0"}, "1", "2"]) + ), }, - selected=selected + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -301,9 +307,11 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_emb() -> N inputs={ "context1": BasedOn(ctx_str_1), "context2": BasedOn(rl_chain.Embed(ctx_str_2)), - "action": ToSelectFrom([{"a": "0", "b": rl_chain.Embed("0")}, "1", rl_chain.Embed("2")]) + "action": ToSelectFrom( + [{"a": "0", "b": rl_chain.Embed("0")}, "1", rl_chain.Embed("2")] + ), }, - selected=selected + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -331,9 +339,15 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_emakeep() inputs={ "context1": BasedOn(ctx_str_1), "context2": BasedOn(rl_chain.EmbedAndKeep(ctx_str_2)), - "action": ToSelectFrom([{"a": "0", "b": rl_chain.EmbedAndKeep("0")}, "1", rl_chain.EmbedAndKeep("2")]) + "action": ToSelectFrom( + [ + {"a": "0", "b": rl_chain.EmbedAndKeep("0")}, + "1", + rl_chain.EmbedAndKeep("2"), + ] + ), }, - selected=selected + selected=selected, ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected) @@ -360,10 +374,7 @@ def test_raw_features_underscored() -> None: ) event = pick_best_chain.PickBestEvent( - inputs={ - "action": ToSelectFrom([str1]), - "context": BasedOn(ctx_str) - } + inputs={"action": ToSelectFrom([str1]), "context": BasedOn(ctx_str)} ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) assert_vw_ex_equals(vw_ex_str, expected_no_embed) @@ -375,7 +386,7 @@ def test_raw_features_underscored() -> None: event = pick_best_chain.PickBestEvent( inputs={ "action": ToSelectFrom(rl_chain.Embed([str1])), - "context": BasedOn(rl_chain.Embed(ctx_str)) + "context": BasedOn(rl_chain.Embed(ctx_str)), } ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event)) @@ -391,7 +402,7 @@ def test_raw_features_underscored() -> None: event = pick_best_chain.PickBestEvent( inputs={ "action": ToSelectFrom(rl_chain.EmbedAndKeep([str1])), - "context": BasedOn(rl_chain.EmbedAndKeep(ctx_str)) + "context": BasedOn(rl_chain.EmbedAndKeep(ctx_str)), } ) vw_ex_str = vw_cb_formatter(*featurizer.featurize(event))