import json
import os
from collections import Counter

import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch import nn
from torch.utils.data import Dataset


lowerCamelCase__  =					{1: (1, 1), 2: (2, 1), 3: (3, 1), 4: (2, 2), 5: (5, 1), 6: (3, 2), 7: (7, 1), 8: (4, 2), 9: (3, 3)}




class 					_UpperCAmelCase			(     nn.Module  ):



 '''simple docstring'''




 def __init__(		self				:   List[str]							, lowercase_				:   str)						->	Any:


   """simple docstring"""







   super().__init__()
   _UpperCamelCase           =    torchvision.models.resnetaaa(pretrained=lowercase_)
   _UpperCamelCase           =    list(model.children())[:-2]
   _UpperCamelCase           =    nn.Sequential(*lowercase_)
   _UpperCamelCase           =    nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds])







 def 		__UpperCAmelCase				(		self				:   Optional[int]							, lowercase_				:   Optional[Any])						->	List[str]:


   """simple docstring"""







   _UpperCamelCase           =    self.pool(self.model(lowercase_))
   _UpperCamelCase           =    torch.flatten(lowercase_							, start_dim=2)
   _UpperCamelCase           =    out.transpose(1							, 2).contiguous()
   return out  # BxNx2048




class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''




 def __init__(		self				:   List[str]							, lowercase_				:   Union[str, Any]							, lowercase_				:   Tuple							, lowercase_				:   Optional[Any]							, lowercase_				:   Any							, lowercase_				:   str)						->	List[Any]:


   """simple docstring"""







   _UpperCamelCase           =    [json.loads(lowercase_) for l in open(lowercase_)]
   _UpperCamelCase           =    os.path.dirname(lowercase_)
   _UpperCamelCase           =    tokenizer
   _UpperCamelCase           =    labels
   _UpperCamelCase           =    len(lowercase_)
   _UpperCamelCase           =    max_seq_length

   _UpperCamelCase           =    transforms







 def __len__(		self				:   Optional[int])						->	int:


   """simple docstring"""







   return len(self.data)







 def __getitem__(		self				:   Tuple							, lowercase_				:   str)						->	Any:


   """simple docstring"""







   _UpperCamelCase           =    torch.LongTensor(self.tokenizer.encode(self.data[index]["text"]							, add_special_tokens=lowercase_))
   _UpperCamelCase      ,						_UpperCamelCase      ,						_UpperCamelCase           =    sentence[0], sentence[1:-1], sentence[-1]
   _UpperCamelCase           =    sentence[: self.max_seq_length]

   _UpperCamelCase           =    torch.zeros(self.n_classes)
   _UpperCamelCase           =    1

   _UpperCamelCase           =    Image.open(os.path.join(self.data_dir							, self.data[index]["img"])).convert("RGB")
   _UpperCamelCase           =    self.transforms(lowercase_)

   return {
       "image_start_token": start_token,
       "image_end_token": end_token,
       "sentence": sentence,
       "image": image,
       "label": label,
   }







 def 		__UpperCAmelCase				(		self				:   int)						->	List[Any]:


   """simple docstring"""







   _UpperCamelCase           =    Counter()
   for row in self.data:
     label_freqs.update(row["label"])
   return label_freqs


def 				lowerCAmelCase__       (			a__ )      ->Dict:




  '''simple docstring'''

  _UpperCamelCase           =    [len(row["sentence"] ) for row in batch]
  _UpperCamelCase      ,						_UpperCamelCase           =    len(a__ ), max(a__ )

  _UpperCamelCase           =    torch.zeros(a__    ,     a__    ,     dtype=torch.long )
  _UpperCamelCase           =    torch.zeros(a__    ,     a__    ,     dtype=torch.long )

  for i_batch, (input_row, length) in enumerate(zip(a__    ,     a__ ) ):
    _UpperCamelCase           =    input_row["sentence"]
    _UpperCamelCase           =    1

  _UpperCamelCase           =    torch.stack([row["image"] for row in batch] )
  _UpperCamelCase           =    torch.stack([row["label"] for row in batch] )
  _UpperCamelCase           =    torch.stack([row["image_start_token"] for row in batch] )
  _UpperCamelCase           =    torch.stack([row["image_end_token"] for row in batch] )

  return text_tensor, mask_tensor, img_tensor, img_start_token, img_end_token, tgt_tensor


def 				lowerCAmelCase__       (			)      ->int:




  '''simple docstring'''

  return [
      "Crime",
      "Drama",
      "Thriller",
      "Action",
      "Comedy",
      "Romance",
      "Documentary",
      "Short",
      "Mystery",
      "History",
      "Family",
      "Adventure",
      "Fantasy",
      "Sci-Fi",
      "Western",
      "Horror",
      "Sport",
      "War",
      "Music",
      "Musical",
      "Animation",
      "Biography",
      "Film-Noir",
  ]







def 				lowerCAmelCase__       (			)      ->int:




  '''simple docstring'''

  return transforms.Compose(
      [
          transforms.Resize(256 ),
          transforms.CenterCrop(224 ),
          transforms.ToTensor(),
          transforms.Normalize(
              mean=[0.46777044, 0.44531429, 0.40661017]    ,     std=[0.12221994, 0.12145835, 0.14380469]    ,     ),
      ] )





# This code is adapted from OpenAI's release
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py

import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import tempfile


def 				lowerCAmelCase__       (			a__    ,     a__    ,     a__    ,     a__ )      ->str:




  '''simple docstring'''

  _UpperCamelCase           =    multiprocessing.Manager()
  _UpperCamelCase           =    manager.list()

  _UpperCamelCase           =    multiprocessing.Process(target=a__    ,     args=(check_program, result, timeout) )
  p.start()
  p.join(timeout=timeout + 1 )
  if p.is_alive():
    p.kill()

  if not result:
    result.append("timed out" )

  return {
      "task_id": task_id,
      "passed": result[0] == "passed",
      "result": result[0],
      "completion_id": completion_id,
  }


def 				lowerCAmelCase__       (			a__    ,     a__    ,     a__ )      ->Union[str, Any]:




  '''simple docstring'''

  with create_tempdir():
    # These system calls are needed when cleaning up tempdir.
    import os
    import shutil

    _UpperCamelCase           =    shutil.rmtree
    _UpperCamelCase           =    os.rmdir
    _UpperCamelCase           =    os.chdir

    # Disable functionalities that can make destructive changes to the test.
    reliability_guard()

    # Run program.
    try:
      _UpperCamelCase           =    {}
      with swallow_io():
        with time_limit(a__ ):
          exec(a__    ,     a__ )
      result.append("passed" )
    except TimeoutException:
      result.append("timed out" )
    except BaseException as e:
      result.append(f'failed: {e}' )

    # Needed for cleaning up.
    _UpperCamelCase           =    rmtree
    _UpperCamelCase           =    rmdir
    _UpperCamelCase           =    chdir


@contextlib.contextmanager
def 				lowerCAmelCase__       (			a__ )      ->Optional[Any]:




  '''simple docstring'''

  def signal_handler(a__    ,     a__ ):
    raise TimeoutException("Timed out!" )

  signal.setitimer(signal.ITIMER_REAL    ,     a__ )
  signal.signal(signal.SIGALRM    ,     a__ )
  try:
    yield
  finally:
    signal.setitimer(signal.ITIMER_REAL    ,     0 )


@contextlib.contextmanager
def 				lowerCAmelCase__       (			)      ->Optional[int]:




  '''simple docstring'''

  _UpperCamelCase           =    WriteOnlyStringIO()
  with contextlib.redirect_stdout(a__ ):
    with contextlib.redirect_stderr(a__ ):
      with redirect_stdin(a__ ):
        yield


@contextlib.contextmanager
def 				lowerCAmelCase__       (			)      ->Optional[int]:




  '''simple docstring'''

  with tempfile.TemporaryDirectory() as dirname:
    with chdir(a__ ):
      yield dirname




class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''




 pass




class 					_UpperCAmelCase			(     io.StringIO  ):



 '''simple docstring'''




 def 		__UpperCAmelCase				(		self				:   List[str]							, *lowercase_				:   str							, **lowercase_				:   List[str])						->	Optional[Any]:


   """simple docstring"""







   raise OSError







 def 		__UpperCAmelCase				(		self				:   Any							, *lowercase_				:   List[Any]							, **lowercase_				:   Optional[int])						->	Optional[int]:


   """simple docstring"""







   raise OSError







 def 		__UpperCAmelCase				(		self				:   int							, *lowercase_				:   int							, **lowercase_				:   str)						->	int:


   """simple docstring"""







   raise OSError







 def 		__UpperCAmelCase				(		self				:   List[Any]							, *lowercase_				:   Tuple							, **lowercase_				:   Tuple)						->	Tuple:


   """simple docstring"""







   return False




class 					_UpperCAmelCase			(     contextlib._RedirectStream  ):  # type: ignore



 '''simple docstring'''




 __A			       =					'''stdin'''


@contextlib.contextmanager
def 				lowerCAmelCase__       (			a__ )      ->str:




  '''simple docstring'''

  if root == ".":
    yield
    return
  _UpperCamelCase           =    os.getcwd()
  os.chdir(a__ )
  try:
    yield
  except BaseException as exc:
    raise exc
  finally:
    os.chdir(a__ )







def 				lowerCAmelCase__       (			a__=None )      ->Optional[int]:




  '''simple docstring'''

  if maximum_memory_bytes is not None:
    import resource

    resource.setrlimit(resource.RLIMIT_AS    ,     (maximum_memory_bytes, maximum_memory_bytes) )
    resource.setrlimit(resource.RLIMIT_DATA    ,     (maximum_memory_bytes, maximum_memory_bytes) )
    if not platform.uname().system == "Darwin":
      resource.setrlimit(resource.RLIMIT_STACK    ,     (maximum_memory_bytes, maximum_memory_bytes) )

  faulthandler.disable()

  import builtins

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None

  import os

  _UpperCamelCase           =    "1"

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None

  import shutil

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None

  import subprocess

  _UpperCamelCase           =    None  # type: ignore

  _UpperCamelCase           =    None

  import sys

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None




