Repository: huggingface/transformers
File: tests/generation/test_utils.py (Python)

Code before:
            fake_model_kwargs = {"foo": "bar"}
            model.generate(input_ids, **fake_model_kwargs)

        # however, valid model_kwargs are accepted
        valid_model_kwargs = {"attention_mask": torch.tensor(np.zeros_like(input_ids))}
        model.generate(input_ids, **valid_model_kwargs)

    def test_custom_logits_processor(self):
        """Tests that custom logits processors can be used in `generate`, and that redundant arguments are caught."""
        bart_tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
        article = """Justin Timberlake and Jessica Biel, welcome to parenthood."""
        bart_model = AutoModelForSeq2SeqLM.from_pretrained("hf-internal-testing/tiny-random-bart")
        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids

        logits_processor = LogitsProcessorList()
        logits_processor.append(MinLengthLogitsProcessor(min_length=10, eos_token_id=0))

        # it should not be allowed to both define `min_length` via kwargs and `logits_processor` list
        logger = logging.get_logger("transformers.generation.utils")
        logger.warning_once.cache_clear()
        with CaptureLogger(logger) as cl:
            bart_model.generate(input_ids, logits_processor=logits_processor, min_length=10)
        self.assertTrue(
            "custom logits processor of type <class 'transformers.generation.logits_process.MinLengthLogitsProcessor'> has been passed to `.generate()`, but it was also created in `.generate()"
            in cl.out
        )

        bart_model.generate(input_ids, logits_processor=logits_processor)

    def test_transition_scores_greedy_search(self):
        """Test that `compute_transition_scores` is working as expected with gready search"""
        articles = ["Justin Timberlake", "Michael Phelps"]
        tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2", padding_side="left")
        tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
        model.generation_config.eos_token_id = None
        input_ids = tokenizer(articles, return_tensors="pt", padding=True).input_ids
        model = model.to(torch_device)
        input_ids = input_ids.to(torch_device)

        outputs = model.generate(
            input_ids=input_ids,
            max_new_tokens=5,
            pad_token_id=tokenizer.eos_token_id,
            return_dict_in_generate=True,
            output_scores=True,
        )

        transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores)
        transition_scores = transition_scores.cpu().numpy()

Review comment (suggestion):
are these emited by a warning? if so it would be more robust to catch the exception on warning instead of relying on a log that may change over time

Code after:
        # arbitrary arguments that will not be used anywhere are also not accepted
        with self.assertRaisesRegex(ValueError, "foo"):
            fake_model_kwargs = {"foo": "bar"}
            model.generate(input_ids, **fake_model_kwargs)

        # however, valid model_kwargs are accepted
        valid_model_kwargs = {"attention_mask": torch.tensor(np.zeros_like(input_ids))}
        model.generate(input_ids, **valid_model_kwargs)

    def test_custom_logits_processor(self):
        """Tests that custom logits processors can be used in `generate`, and that redundant arguments are caught."""
        bart_tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
        article = """Justin Timberlake and Jessica Biel, welcome to parenthood."""
        bart_model = AutoModelForSeq2SeqLM.from_pretrained("hf-internal-testing/tiny-random-bart")
        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids

        logits_processor = LogitsProcessorList()
        logits_processor.append(MinLengthLogitsProcessor(min_length=10, eos_token_id=0))

        # it should not be allowed to both define `min_length` via kwargs and `logits_processor` list
        logger = logging.get_logger("transformers.generation.utils")
        logger.warning_once.cache_clear()
        with CaptureLogger(logger) as cl:
            bart_model.generate(input_ids, logits_processor=logits_processor, min_length=10)
        self.assertTrue(
            "custom logits processor of type <class 'transformers.generation.logits_process.MinLengthLogitsProcessor'> has been passed to `.generate()`, but it was also created in `.generate()"
            in cl.out
        )

        bart_model.generate(input_ids, logits_processor=logits_processor)

    def test_transition_scores_greedy_search(self):
        """Test that `compute_transition_scores` is working as expected with gready search"""
        articles = ["Justin Timberlake", "Michael Phelps"]
        tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2", padding_side="left")
        tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
        model.generation_config.eos_token_id = None
        input_ids = tokenizer(articles, return_tensors="pt", padding=True).input_ids
        model = model.to(torch_device)
        input_ids = input_ids.to(torch_device)

        outputs = model.generate(
            input_ids=input_ids,
            max_new_tokens=5,
            pad_token_id=tokenizer.eos_token_id,
            return_dict_in_generate=True,
            output_scores=True,
        )