add targets arg to fill-mask pipeline (#6239)
* add targets arg to fill-mask pipeline * add tests and more error handling * quality * update docstring
This commit is contained in:
@@ -41,6 +41,23 @@ expected_fill_mask_result = [
|
||||
],
|
||||
]
|
||||
|
||||
expected_fill_mask_target_result = [
|
||||
[
|
||||
{
|
||||
"sequence": "<s>My name is Patrick</s>",
|
||||
"score": 0.004992353264242411,
|
||||
"token": 3499,
|
||||
"token_str": "ĠPatrick",
|
||||
},
|
||||
{
|
||||
"sequence": "<s>My name is Clara</s>",
|
||||
"score": 0.00019297805556561798,
|
||||
"token": 13606,
|
||||
"token_str": "ĠClara",
|
||||
},
|
||||
]
|
||||
]
|
||||
|
||||
SUMMARIZATION_KWARGS = dict(num_beams=2, min_length=2, max_length=5)
|
||||
|
||||
|
||||
@@ -139,7 +156,7 @@ class MonoColumnInputTestCase(unittest.TestCase):
|
||||
for key in output_keys:
|
||||
self.assertIn(key, mono_result[0])
|
||||
|
||||
multi_result = [nlp(input) for input in valid_inputs]
|
||||
multi_result = [nlp(input, **kwargs) for input in valid_inputs]
|
||||
self.assertIsInstance(multi_result, list)
|
||||
self.assertIsInstance(multi_result[0], (dict, list))
|
||||
|
||||
@@ -219,6 +236,34 @@ class MonoColumnInputTestCase(unittest.TestCase):
|
||||
nlp, valid_inputs, mandatory_keys, invalid_inputs, expected_check_keys=["sequence"]
|
||||
)
|
||||
|
||||
@require_torch
|
||||
def test_torch_fill_mask_with_targets(self):
|
||||
valid_inputs = ["My name is <mask>"]
|
||||
valid_targets = [[" Teven", " Patrick", " Clara"], [" Sam"]]
|
||||
invalid_targets = [[], [""], ""]
|
||||
for model_name in FILL_MASK_FINETUNED_MODELS:
|
||||
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt")
|
||||
for targets in valid_targets:
|
||||
outputs = nlp(valid_inputs, targets=targets)
|
||||
self.assertIsInstance(outputs, list)
|
||||
self.assertEqual(len(outputs), len(targets))
|
||||
for targets in invalid_targets:
|
||||
self.assertRaises(ValueError, nlp, valid_inputs, targets=targets)
|
||||
|
||||
@require_tf
|
||||
def test_tf_fill_mask_with_targets(self):
|
||||
valid_inputs = ["My name is <mask>"]
|
||||
valid_targets = [[" Teven", " Patrick", " Clara"], [" Sam"]]
|
||||
invalid_targets = [[], [""], ""]
|
||||
for model_name in FILL_MASK_FINETUNED_MODELS:
|
||||
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf")
|
||||
for targets in valid_targets:
|
||||
outputs = nlp(valid_inputs, targets=targets)
|
||||
self.assertIsInstance(outputs, list)
|
||||
self.assertEqual(len(outputs), len(targets))
|
||||
for targets in invalid_targets:
|
||||
self.assertRaises(ValueError, nlp, valid_inputs, targets=targets)
|
||||
|
||||
@require_torch
|
||||
@slow
|
||||
def test_torch_fill_mask_results(self):
|
||||
@@ -227,6 +272,7 @@ class MonoColumnInputTestCase(unittest.TestCase):
|
||||
"My name is <mask>",
|
||||
"The largest city in France is <mask>",
|
||||
]
|
||||
valid_targets = [" Patrick", " Clara"]
|
||||
for model_name in LARGE_FILL_MASK_FINETUNED_MODELS:
|
||||
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt", topk=2,)
|
||||
self._test_mono_column_pipeline(
|
||||
@@ -236,6 +282,14 @@ class MonoColumnInputTestCase(unittest.TestCase):
|
||||
expected_multi_result=expected_fill_mask_result,
|
||||
expected_check_keys=["sequence"],
|
||||
)
|
||||
self._test_mono_column_pipeline(
|
||||
nlp,
|
||||
valid_inputs[:1],
|
||||
mandatory_keys,
|
||||
expected_multi_result=expected_fill_mask_target_result,
|
||||
expected_check_keys=["sequence"],
|
||||
targets=valid_targets,
|
||||
)
|
||||
|
||||
@require_tf
|
||||
@slow
|
||||
@@ -245,6 +299,7 @@ class MonoColumnInputTestCase(unittest.TestCase):
|
||||
"My name is <mask>",
|
||||
"The largest city in France is <mask>",
|
||||
]
|
||||
valid_targets = [" Patrick", " Clara"]
|
||||
for model_name in LARGE_FILL_MASK_FINETUNED_MODELS:
|
||||
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf", topk=2)
|
||||
self._test_mono_column_pipeline(
|
||||
@@ -254,6 +309,14 @@ class MonoColumnInputTestCase(unittest.TestCase):
|
||||
expected_multi_result=expected_fill_mask_result,
|
||||
expected_check_keys=["sequence"],
|
||||
)
|
||||
self._test_mono_column_pipeline(
|
||||
nlp,
|
||||
valid_inputs[:1],
|
||||
mandatory_keys,
|
||||
expected_multi_result=expected_fill_mask_target_result,
|
||||
expected_check_keys=["sequence"],
|
||||
targets=valid_targets,
|
||||
)
|
||||
|
||||
@require_torch
|
||||
def test_torch_summarization(self):
|
||||
|
||||
Reference in New Issue
Block a user