



"""simple docstring"""
import torch
from torch import nn


class    SCREAMING_SNAKE_CASE__ (							nn.Module   ):

    def __init__(				self			:				Dict ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				List[Any] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				List[str]=1 ,  lowerCAmelCase_			:				int=False):







      """simple docstring"""







      super().__init__()

      lowercase_								= n_token
      lowercase_								= d_embed
      lowercase_								= d_proj

      lowercase_								= cutoffs + [n_token]
      lowercase_								= [0] + self.cutoffs
      lowercase_								= div_val

      lowercase_								= self.cutoffs[0]
      lowercase_								= len(self.cutoffs) - 1
      lowercase_								= self.shortlist_size + self.n_clusters

      if self.n_clusters > 0:
        lowercase_								= nn.Parameter(torch.zeros(self.n_clusters ,  self.d_embed))
        lowercase_								= nn.Parameter(torch.zeros(self.n_clusters))

      lowercase_								= nn.ModuleList()
      lowercase_								= nn.ParameterList()

      if div_val == 1:
        for i in range(len(self.cutoffs)):
          if d_proj != d_embed:
            self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCAmelCase_ ,  lowerCAmelCase_)))
          else:
            self.out_projs.append(lowerCAmelCase_)

        self.out_layers.append(nn.Linear(lowerCAmelCase_ ,  lowerCAmelCase_))
      else:
        for i in range(len(self.cutoffs)):
          lowercase_					,    lowercase_								= self.cutoff_ends[i], self.cutoff_ends[i + 1]
          lowercase_								= d_embed // (div_val**i)

          self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCAmelCase_ ,  lowerCAmelCase_)))

          self.out_layers.append(nn.Linear(lowerCAmelCase_ ,  r_idx - l_idx))

      lowercase_								= keep_order

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Dict):







      """simple docstring"""







      if proj is None:
        lowercase_								= nn.functional.linear(lowerCAmelCase_ ,  lowerCAmelCase_ ,  bias=lowerCAmelCase_)
      else:
        # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
        lowercase_								= nn.functional.linear(lowerCAmelCase_ ,  proj.t().contiguous())
        lowercase_								= nn.functional.linear(lowerCAmelCase_ ,  lowerCAmelCase_ ,  bias=lowerCAmelCase_)
        # else:
        #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
        #     if bias is not None:
        #         logit = logit + bias

      return logit

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Optional[Any]=None ,  lowerCAmelCase_			:				List[Any]=False):







      """simple docstring"""







      if labels is not None:
        # Shift so that tokens < n predict n
        lowercase_								= hidden[..., :-1, :].contiguous()
        lowercase_								= labels[..., 1:].contiguous()
        lowercase_								= hidden.view(-1 ,  hidden.size(-1))
        lowercase_								= labels.view(-1)
        if hidden.size(0) != labels.size(0):
          raise RuntimeError("""Input and labels should have the same size in the batch dimension.""")
      else:
        lowercase_								= hidden.view(-1 ,  hidden.size(-1))

      if self.n_clusters == 0:
        lowercase_								= self._compute_logit(lowerCAmelCase_ ,  self.out_layers[0].weight ,  self.out_layers[0].bias ,  self.out_projs[0])
        if labels is not None:
          lowercase_								= labels != -1_0_0
          lowercase_								= torch.zeros_like(lowerCAmelCase_ ,  dtype=hidden.dtype ,  device=hidden.device)
          lowercase_								= (
              -nn.functional.log_softmax(lowerCAmelCase_ ,  dim=-1)[mask].gather(1 ,  labels[mask].unsqueeze(1)).squeeze(1)
          )
        else:
          lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=-1)
      else:
        # construct weights and biases
        lowercase_					,    lowercase_								= [], []
        for i in range(len(self.cutoffs)):
          if self.div_val == 1:
            lowercase_					,    lowercase_								= self.cutoff_ends[i], self.cutoff_ends[i + 1]
            lowercase_								= self.out_layers[0].weight[l_idx:r_idx]
            lowercase_								= self.out_layers[0].bias[l_idx:r_idx]
          else:
            lowercase_								= self.out_layers[i].weight
            lowercase_								= self.out_layers[i].bias

          if i == 0:
            lowercase_								= torch.cat([weight_i, self.cluster_weight] ,  dim=0)
            lowercase_								= torch.cat([bias_i, self.cluster_bias] ,  dim=0)

          weights.append(lowerCAmelCase_)
          biases.append(lowerCAmelCase_)

        lowercase_					,    lowercase_					,    lowercase_								= weights[0], biases[0], self.out_projs[0]

        lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
        lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=1)

        if labels is None:
          lowercase_								= hidden.new_empty((head_logit.size(0), self.n_token))
        else:
          lowercase_								= torch.zeros_like(lowerCAmelCase_ ,  dtype=hidden.dtype ,  device=hidden.device)

        lowercase_								= 0
        lowercase_								= [0] + self.cutoffs
        for i in range(len(lowerCAmelCase_) - 1):
          lowercase_					,    lowercase_								= cutoff_values[i], cutoff_values[i + 1]

          if labels is not None:
            lowercase_								= (labels >= l_idx) & (labels < r_idx)
            lowercase_								= mask_i.nonzero().squeeze()

            if indices_i.numel() == 0:
              continue

            lowercase_								= labels.index_select(0 ,  lowerCAmelCase_) - l_idx
            lowercase_								= head_logprob.index_select(0 ,  lowerCAmelCase_)
            lowercase_								= hidden.index_select(0 ,  lowerCAmelCase_)
          else:
            lowercase_								= hidden

          if i == 0:
            if labels is not None:
              lowercase_								= head_logprob_i.gather(1 ,  target_i[:, None]).squeeze(1)
            else:
              lowercase_								= head_logprob[:, : self.cutoffs[0]]
          else:
            lowercase_					,    lowercase_					,    lowercase_								= weights[i], biases[i], self.out_projs[i]

            lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
            lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=1)
            lowercase_								= self.cutoffs[0] + i - 1  # No probability for the head cluster
            if labels is not None:
              lowercase_								= head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                  1 ,  target_i[:, None]).squeeze(1)
            else:
              lowercase_								= head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
              lowercase_								= logprob_i

          if labels is not None:
            if (hasattr(self ,  """keep_order""") and self.keep_order) or keep_order:
              out.index_copy_(0 ,  lowerCAmelCase_ ,  -logprob_i)
            else:
              out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
            offset += logprob_i.size(0)

      return out




    def 					_UpperCAmelCase	(				self			:				List[Any] ,  lowerCAmelCase_			:				str):







      """simple docstring"""







      if self.n_clusters == 0:
        lowercase_								= self._compute_logit(lowerCAmelCase_ ,  self.out_layers[0].weight ,  self.out_layers[0].bias ,  self.out_projs[0])
        return nn.functional.log_softmax(lowerCAmelCase_ ,  dim=-1)
      else:
        # construct weights and biases
        lowercase_					,    lowercase_								= [], []
        for i in range(len(self.cutoffs)):
          if self.div_val == 1:
            lowercase_					,    lowercase_								= self.cutoff_ends[i], self.cutoff_ends[i + 1]
            lowercase_								= self.out_layers[0].weight[l_idx:r_idx]
            lowercase_								= self.out_layers[0].bias[l_idx:r_idx]
          else:
            lowercase_								= self.out_layers[i].weight
            lowercase_								= self.out_layers[i].bias

          if i == 0:
            lowercase_								= torch.cat([weight_i, self.cluster_weight] ,  dim=0)
            lowercase_								= torch.cat([bias_i, self.cluster_bias] ,  dim=0)

          weights.append(lowerCAmelCase_)
          biases.append(lowerCAmelCase_)

        lowercase_					,    lowercase_					,    lowercase_								= weights[0], biases[0], self.out_projs[0]
        lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)

        lowercase_								= hidden.new_empty((head_logit.size(0), self.n_token))
        lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=1)

        lowercase_								= [0] + self.cutoffs
        for i in range(len(lowerCAmelCase_) - 1):
          lowercase_					,    lowercase_								= cutoff_values[i], cutoff_values[i + 1]

          if i == 0:
            lowercase_								= head_logprob[:, : self.cutoffs[0]]
          else:
            lowercase_					,    lowercase_					,    lowercase_								= weights[i], biases[i], self.out_projs[i]

            lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
            lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=1)

            lowercase_								= head_logprob[:, -i] + tail_logprob_i
            lowercase_								= logprob_i

        return out






"""simple docstring"""
import unittest

from transformers import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING, is_vision_available, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
       from PIL import Image
else:


       class    SCREAMING_SNAKE_CASE__ :

           @staticmethod
           def 					_UpperCAmelCase	(				*lowerCAmelCase_			:				Tuple ,  **lowerCAmelCase_			:				Tuple):







             """simple docstring"""







             pass


@is_pipeline_test
@require_vision
@require_torch
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):
    lowercase__      =   MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

    def 					_UpperCAmelCase	(				self			:				Optional[Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      lowercase_								= pipeline(
          """zero-shot-object-detection""" ,  model="""hf-internal-testing/tiny-random-owlvit-object-detection""")

      lowercase_								= [
          {
              """image""": """./tests/fixtures/tests_samples/COCO/000000039769.png""",
              """candidate_labels""": ["""cat""", """remote""", """couch"""],
          }
      ]
      return object_detector, examples

    def 					_UpperCAmelCase	(				self			:				Optional[Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= object_detector(examples[0] ,  threshold=0.0)

      lowercase_								= len(lowerCAmelCase_)
      self.assertGreater(lowerCAmelCase_ ,  0)
      self.assertEqual(
          lowerCAmelCase_ ,  [
              {
                  """score""": ANY(lowerCAmelCase_),
                  """label""": ANY(lowerCAmelCase_),
                  """box""": {"""xmin""": ANY(lowerCAmelCase_), """ymin""": ANY(lowerCAmelCase_), """xmax""": ANY(lowerCAmelCase_), """ymax""": ANY(lowerCAmelCase_)},
              }
              for i in range(lowerCAmelCase_)
          ] ,  )

    @require_tf
    @unittest.skip("""Zero Shot Object Detection not implemented in TF""")
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      pass

    @require_torch
    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_								= pipeline(
          """zero-shot-object-detection""" ,  model="""hf-internal-testing/tiny-random-owlvit-object-detection""")

      lowercase_								= object_detector(
          """./tests/fixtures/tests_samples/COCO/000000039769.png""" ,  candidate_labels=["""cat""", """remote""", """couch"""] ,  threshold=0.64 ,  )

      self.assertEqual(
          nested_simplify(lowerCAmelCase_ ,  decimals=4) ,  [
              {"""score""": 0.7_235, """label""": """cat""", """box""": {"""xmin""": 2_0_4, """ymin""": 1_6_7, """xmax""": 2_3_2, """ymax""": 1_9_0}},
              {"""score""": 0.7_218, """label""": """remote""", """box""": {"""xmin""": 2_0_4, """ymin""": 1_6_7, """xmax""": 2_3_2, """ymax""": 1_9_0}},
              {"""score""": 0.7_184, """label""": """couch""", """box""": {"""xmin""": 2_0_4, """ymin""": 1_6_7, """xmax""": 2_3_2, """ymax""": 1_9_0}},
              {"""score""": 0.6_748, """label""": """remote""", """box""": {"""xmin""": 5_7_1, """ymin""": 8_3, """xmax""": 5_9_8, """ymax""": 1_0_3}},
              {"""score""": 0.6_656, """label""": """cat""", """box""": {"""xmin""": 5_7_1, """ymin""": 8_3, """xmax""": 5_9_8, """ymax""": 1_0_3}},
              {"""score""": 0.6_614, """label""": """couch""", """box""": {"""xmin""": 5_7_1, """ymin""": 8_3, """xmax""": 5_9_8, """ymax""": 1_0_3}},
              {"""score""": 0.6_456, """label""": """remote""", """box""": {"""xmin""": 4_9_4, """ymin""": 1_0_5, """xmax""": 5_2_1, """ymax""": 1_2_7}},
              {"""score""": 0.642, """label""": """remote""", """box""": {"""xmin""": 6_7, """ymin""": 2_7_4, """xmax""": 9_3, """ymax""": 2_9_7}},
              {"""score""": 0.6_419, """label""": """cat""", """box""": {"""xmin""": 4_9_4, """ymin""": 1_0_5, """xmax""": 5_2_1, """ymax""": 1_2_7}},
          ] ,  )

      lowercase_								= object_detector(
          [
              {
                  """image""": """./tests/fixtures/tests_samples/COCO/000000039769.png""",
                  """candidate_labels""": ["""cat""", """remote""", """couch"""],
              }
          ] ,  threshold=0.64 ,  )

      self.assertEqual(
          nested_simplify(lowerCAmelCase_ ,  decimals=4) ,  [
              [
                  {"""score""": 0.7_235, """label""": """cat""", """box""": {"""xmin""": 2_0_4, """ymin""": 1_6_7, """xmax""": 2_3_2, """ymax""": 1_9_0}},
                  {"""score""": 0.7_218, """label""": """remote""", """box""": {"""xmin""": 2_0_4, """ymin""": 1_6_7, """xmax""": 2_3_2, """ymax""": 1_9_0}},
                  {"""score""": 0.7_184, """label""": """couch""", """box""": {"""xmin""": 2_0_4, """ymin""": 1_6_7, """xmax""": 2_3_2, """ymax""": 1_9_0}},
                  {"""score""": 0.6_748, """label""": """remote""", """box""": {"""xmin""": 5_7_1, """ymin""": 8_3, """xmax""": 5_9_8, """ymax""": 1_0_3}},
                  {"""score""": 0.6_656, """label""": """cat""", """box""": {"""xmin""": 5_7_1, """ymin""": 8_3, """xmax""": 5_9_8, """ymax""": 1_0_3}},
                  {"""score""": 0.6_614, """label""": """couch""", """box""": {"""xmin""": 5_7_1, """ymin""": 8_3, """xmax""": 5_9_8, """ymax""": 1_0_3}},
                  {"""score""": 0.6_456, """label""": """remote""", """box""": {"""xmin""": 4_9_4, """ymin""": 1_0_5, """xmax""": 5_2_1, """ymax""": 1_2_7}},
                  {"""score""": 0.642, """label""": """remote""", """box""": {"""xmin""": 6_7, """ymin""": 2_7_4, """xmax""": 9_3, """ymax""": 2_9_7}},
                  {"""score""": 0.6_419, """label""": """cat""", """box""": {"""xmin""": 4_9_4, """ymin""": 1_0_5, """xmax""": 5_2_1, """ymax""": 1_2_7}},
              ]
          ] ,  )

    @require_torch
    @slow
    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= pipeline("""zero-shot-object-detection""")

      lowercase_								= object_detector(
          """http://images.cocodataset.org/val2017/000000039769.jpg""" ,  candidate_labels=["""cat""", """remote""", """couch"""] ,  )
      self.assertEqual(
          nested_simplify(lowerCAmelCase_ ,  decimals=4) ,  [
              {"""score""": 0.2_868, """label""": """cat""", """box""": {"""xmin""": 3_2_4, """ymin""": 2_0, """xmax""": 6_4_0, """ymax""": 3_7_3}},
              {"""score""": 0.277, """label""": """remote""", """box""": {"""xmin""": 4_0, """ymin""": 7_2, """xmax""": 1_7_7, """ymax""": 1_1_5}},
              {"""score""": 0.2_537, """label""": """cat""", """box""": {"""xmin""": 1, """ymin""": 5_5, """xmax""": 3_1_5, """ymax""": 4_7_2}},
              {"""score""": 0.1_474, """label""": """remote""", """box""": {"""xmin""": 3_3_5, """ymin""": 7_4, """xmax""": 3_7_1, """ymax""": 1_8_7}},
              {"""score""": 0.1_208, """label""": """couch""", """box""": {"""xmin""": 4, """ymin""": 0, """xmax""": 6_4_2, """ymax""": 4_7_6}},
          ] ,  )

      lowercase_								= object_detector(
          [
              {
                  """image""": """http://images.cocodataset.org/val2017/000000039769.jpg""",
                  """candidate_labels""": ["""cat""", """remote""", """couch"""],
              },
              {
                  """image""": """http://images.cocodataset.org/val2017/000000039769.jpg""",
                  """candidate_labels""": ["""cat""", """remote""", """couch"""],
              },
          ] ,  )
      self.assertEqual(
          nested_simplify(lowerCAmelCase_ ,  decimals=4) ,  [
              [
                  {"""score""": 0.2_868, """label""": """cat""", """box""": {"""xmin""": 3_2_4, """ymin""": 2_0, """xmax""": 6_4_0, """ymax""": 3_7_3}},
                  {"""score""": 0.277, """label""": """remote""", """box""": {"""xmin""": 4_0, """ymin""": 7_2, """xmax""": 1_7_7, """ymax""": 1_1_5}},
                  {"""score""": 0.2_537, """label""": """cat""", """box""": {"""xmin""": 1, """ymin""": 5_5, """xmax""": 3_1_5, """ymax""": 4_7_2}},
                  {"""score""": 0.1_474, """label""": """remote""", """box""": {"""xmin""": 3_3_5, """ymin""": 7_4, """xmax""": 3_7_1, """ymax""": 1_8_7}},
                  {"""score""": 0.1_208, """label""": """couch""", """box""": {"""xmin""": 4, """ymin""": 0, """xmax""": 6_4_2, """ymax""": 4_7_6}},
              ],
              [
                  {"""score""": 0.2_868, """label""": """cat""", """box""": {"""xmin""": 3_2_4, """ymin""": 2_0, """xmax""": 6_4_0, """ymax""": 3_7_3}},
                  {"""score""": 0.277, """label""": """remote""", """box""": {"""xmin""": 4_0, """ymin""": 7_2, """xmax""": 1_7_7, """ymax""": 1_1_5}},
                  {"""score""": 0.2_537, """label""": """cat""", """box""": {"""xmin""": 1, """ymin""": 5_5, """xmax""": 3_1_5, """ymax""": 4_7_2}},
                  {"""score""": 0.1_474, """label""": """remote""", """box""": {"""xmin""": 3_3_5, """ymin""": 7_4, """xmax""": 3_7_1, """ymax""": 1_8_7}},
                  {"""score""": 0.1_208, """label""": """couch""", """box""": {"""xmin""": 4, """ymin""": 0, """xmax""": 6_4_2, """ymax""": 4_7_6}},
              ],
          ] ,  )

    @require_tf
    @unittest.skip("""Zero Shot Object Detection not implemented in TF""")
    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      pass

    @require_torch
    @slow
    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      lowercase_								= 0.2
      lowercase_								= pipeline("""zero-shot-object-detection""")

      lowercase_								= object_detector(
          """http://images.cocodataset.org/val2017/000000039769.jpg""" ,  candidate_labels=["""cat""", """remote""", """couch"""] ,  threshold=lowerCAmelCase_ ,  )
      self.assertEqual(
          nested_simplify(lowerCAmelCase_ ,  decimals=4) ,  [
              {"""score""": 0.2_868, """label""": """cat""", """box""": {"""xmin""": 3_2_4, """ymin""": 2_0, """xmax""": 6_4_0, """ymax""": 3_7_3}},
              {"""score""": 0.277, """label""": """remote""", """box""": {"""xmin""": 4_0, """ymin""": 7_2, """xmax""": 1_7_7, """ymax""": 1_1_5}},
              {"""score""": 0.2_537, """label""": """cat""", """box""": {"""xmin""": 1, """ymin""": 5_5, """xmax""": 3_1_5, """ymax""": 4_7_2}},
          ] ,  )




    @require_torch
    @slow
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      lowercase_								= 2
      lowercase_								= pipeline("""zero-shot-object-detection""")

      lowercase_								= object_detector(
          """http://images.cocodataset.org/val2017/000000039769.jpg""" ,  candidate_labels=["""cat""", """remote""", """couch"""] ,  top_k=lowerCAmelCase_ ,  )
      self.assertEqual(
          nested_simplify(lowerCAmelCase_ ,  decimals=4) ,  [
              {"""score""": 0.2_868, """label""": """cat""", """box""": {"""xmin""": 3_2_4, """ymin""": 2_0, """xmax""": 6_4_0, """ymax""": 3_7_3}},
              {"""score""": 0.277, """label""": """remote""", """box""": {"""xmin""": 4_0, """ymin""": 7_2, """xmax""": 1_7_7, """ymax""": 1_1_5}},
          ] ,  )

