

"""simple docstring"""
import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
    import torch

    from transformers import AutoModelForCausalLM

@require_torch
class       lowerCamelCase	(       unittest.TestCase   ):


  '''simple docstring'''







  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Any							=    AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
        UpperCAmelCase__			:							Dict							=    AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							int							=    -1

        UpperCAmelCase__			:							str							=    ids_tensor((1, 5)							,       vocab_size=model.config.vocab_size	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Tuple							=    model.generate(_lowerCamelCase							,       max_new_tokens=10							,       do_sample=_lowerCamelCase	)
        UpperCAmelCase__			:							Optional[int]							=    tokenizer.decode(greedy_ids[0]	)

        with CaptureStdout() as cs:
              UpperCAmelCase__			:							Tuple							=    TextStreamer(_lowerCamelCase	)
              model.generate(_lowerCamelCase							,       max_new_tokens=10							,       do_sample=_lowerCamelCase							,       streamer=_lowerCamelCase	)
        # The greedy text should be printed to stdout, except for the final "\n" in the streamer
        UpperCAmelCase__			:							List[str]							=    cs.out[:-1]

        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[Any]							=    AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
        UpperCAmelCase__			:							List[str]							=    AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Tuple							=    -1

        UpperCAmelCase__			:							Tuple							=    ids_tensor((1, 5)							,       vocab_size=model.config.vocab_size	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Union[str, Any]							=    model.generate(_lowerCamelCase							,       max_new_tokens=10							,       do_sample=_lowerCamelCase	)
        UpperCAmelCase__			:							Tuple							=    tokenizer.decode(greedy_ids[0]	)

        UpperCAmelCase__			:							Union[str, Any]							=    TextIteratorStreamer(_lowerCamelCase	)
        UpperCAmelCase__			:							Tuple							=    {"""input_ids""": input_ids, """max_new_tokens""": 10, """do_sample""": False, """streamer""": streamer}
        UpperCAmelCase__			:							List[Any]							=    Thread(target=model.generate							,       kwargs=_lowerCamelCase	)
        thread.start()
        UpperCAmelCase__			:							Optional[int]							=    """"""
        for new_text in streamer:
              streamer_text += new_text

        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Dict							=    AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
        UpperCAmelCase__			:							Optional[Any]							=    AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Optional[Any]							=    -1

        UpperCAmelCase__			:							List[Any]							=    ids_tensor((1, 5)							,       vocab_size=model.config.vocab_size	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							int							=    model.generate(_lowerCamelCase							,       max_new_tokens=10							,       do_sample=_lowerCamelCase	)
        UpperCAmelCase__			:							int							=    greedy_ids[:, input_ids.shape[1] :]
        UpperCAmelCase__			:							Optional[Any]							=    tokenizer.decode(new_greedy_ids[0]	)

        with CaptureStdout() as cs:
              UpperCAmelCase__			:							Dict							=    TextStreamer(_lowerCamelCase							,       skip_prompt=_lowerCamelCase	)
              model.generate(_lowerCamelCase							,       max_new_tokens=10							,       do_sample=_lowerCamelCase							,       streamer=_lowerCamelCase	)
        # The greedy text should be printed to stdout, except for the final "\n" in the streamer
        UpperCAmelCase__			:							str							=    cs.out[:-1]

        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[Any]							=    AutoTokenizer.from_pretrained("""distilgpt2"""	)
        UpperCAmelCase__			:							List[str]							=    AutoModelForCausalLM.from_pretrained("""distilgpt2"""	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							List[str]							=    -1

        UpperCAmelCase__			:							List[str]							=    torch.ones((1, 5)							,       device=_lowerCamelCase	).long() * model.config.bos_token_id
        with CaptureStdout() as cs:
              UpperCAmelCase__			:							Any							=    TextStreamer(_lowerCamelCase							,       skip_special_tokens=_lowerCamelCase	)
              model.generate(_lowerCamelCase							,       max_new_tokens=1							,       do_sample=_lowerCamelCase							,       streamer=_lowerCamelCase	)

        # The prompt contains a special token, so the streamer should not print it. As such, the output text, when
        # re-tokenized, must only contain one token
        UpperCAmelCase__			:							Optional[Any]							=    cs.out[:-1]  # Remove the final "\n"
        UpperCAmelCase__			:							int							=    tokenizer(_lowerCamelCase							,       return_tensors="""pt"""	)
        self.assertEqual(streamer_text_tokenized.input_ids.shape							,       (1, 1)	)




  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
        UpperCAmelCase__			:							List[str]							=    AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							str							=    -1

        UpperCAmelCase__			:							List[str]							=    ids_tensor((1, 5)							,       vocab_size=model.config.vocab_size	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Optional[Any]							=    TextIteratorStreamer(_lowerCamelCase							,       timeout=0.001	)
        UpperCAmelCase__			:							Union[str, Any]							=    {"""input_ids""": input_ids, """max_new_tokens""": 10, """do_sample""": False, """streamer""": streamer}
        UpperCAmelCase__			:							Dict							=    Thread(target=model.generate							,       kwargs=_lowerCamelCase	)
        thread.start()

        # The streamer will timeout after 0.001 seconds, so an exception will be raised
        with self.assertRaises(_lowerCamelCase	):
              UpperCAmelCase__			:							str							=    """"""
              for new_text in streamer:
                    streamer_text += new_text




"""simple docstring"""
import unittest

import numpy as np
import torch

from .utils_summarization import build_mask, compute_token_type_ids, process_story, truncate_or_pad

class       lowerCamelCase	(       unittest.TestCase   ):


  '''simple docstring'''







  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    10



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    [1, 2, 3, 4]
        UpperCAmelCase__			:							Tuple							=    [1, 2, 3, 4, 0, 0, 0, 0, 0, 0]
        self.assertEqual(truncate_or_pad(_lowerCamelCase							,       self.block_size							,       0	)							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[int]							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        UpperCAmelCase__			:							Tuple							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        self.assertEqual(truncate_or_pad(_lowerCamelCase							,       self.block_size							,       0	)							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Tuple							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
        UpperCAmelCase__			:							Dict							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        self.assertEqual(truncate_or_pad(_lowerCamelCase							,       self.block_size							,       0	)							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Any							=    """It was the year of Our Lord one thousand seven hundred and
        seventy-five.\n\nSpiritual revelations were conceded to England at that
        favoured period, as at this."""
        UpperCAmelCase__					,		UpperCAmelCase__			:							Any							=    process_story(_lowerCamelCase	)
        self.assertEqual(_lowerCamelCase							,       []	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							str							=    """"""
        UpperCAmelCase__					,		UpperCAmelCase__			:							List[Any]							=    process_story(_lowerCamelCase	)
        self.assertEqual(_lowerCamelCase							,       []	)
        self.assertEqual(_lowerCamelCase							,       []	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[Any]							=    (
            """It was the year of Our Lord one thousand seven hundred and """
            """seventy-five\n\nSpiritual revelations were conceded to England """
            """at that favoured period, as at this.\n@highlight\n\nIt was the best of times"""
        )
        UpperCAmelCase__					,		UpperCAmelCase__			:							Optional[Any]							=    process_story(_lowerCamelCase	)

        UpperCAmelCase__			:							List[Any]							=    [
            """It was the year of Our Lord one thousand seven hundred and seventy-five.""",
            """Spiritual revelations were conceded to England at that favoured period, as at this.""",
        ]
        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)

        UpperCAmelCase__			:							List[Any]							=    ["""It was the best of times."""]
        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							str							=    torch.tensor([1, 2, 3, 4]	)
        UpperCAmelCase__			:							str							=    torch.tensor([1, 1, 1, 1]	)
        np.testing.assert_array_equal(build_mask(_lowerCamelCase							,       0	).numpy()							,       expected.numpy()	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    torch.tensor([1, 2, 3, 4, 23, 23, 23]	)
        UpperCAmelCase__			:							Union[str, Any]							=    torch.tensor([1, 1, 1, 1, 0, 0, 0]	)
        np.testing.assert_array_equal(build_mask(_lowerCamelCase							,       23	).numpy()							,       expected.numpy()	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[Any]							=    torch.tensor([8, 2, 3, 4, 1, 1, 1]	)
        UpperCAmelCase__			:							Any							=    torch.tensor([1, 1, 1, 1, 0, 0, 0]	)
        np.testing.assert_array_equal(build_mask(_lowerCamelCase							,       1	).numpy()							,       expected.numpy()	)




  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Any							=    101
        UpperCAmelCase__			:							Dict							=    torch.tensor([[1, 2, 3, 4, 5, 6], [1, 2, 3, 101, 5, 6], [1, 101, 3, 4, 101, 6]]	)
        UpperCAmelCase__			:							Optional[Any]							=    torch.tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0], [1, 0, 0, 0, 1, 1]]	)

        UpperCAmelCase__			:							int							=    compute_token_type_ids(_lowerCamelCase							,       _lowerCamelCase	)
        np.testing.assert_array_equal(_lowerCamelCase							,       _lowerCamelCase	)

