






import argparse
import re

import numpy as np
import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SamConfig,
    SamImageProcessor,
    SamModel,
    SamProcessor,
    SamVisionConfig,
)


lowerCamelCase__       		=  {
    """iou_prediction_head.layers.0""": """iou_prediction_head.proj_in""",
    """iou_prediction_head.layers.1""": """iou_prediction_head.layers.0""",
    """iou_prediction_head.layers.2""": """iou_prediction_head.proj_out""",
    """mask_decoder.output_upscaling.0""": """mask_decoder.upscale_conv1""",
    """mask_decoder.output_upscaling.1""": """mask_decoder.upscale_layer_norm""",
    """mask_decoder.output_upscaling.3""": """mask_decoder.upscale_conv2""",
    """mask_downscaling.0""": """mask_embed.conv1""",
    """mask_downscaling.1""": """mask_embed.layer_norm1""",
    """mask_downscaling.3""": """mask_embed.conv2""",
    """mask_downscaling.4""": """mask_embed.layer_norm2""",
    """mask_downscaling.6""": """mask_embed.conv3""",
    """point_embeddings""": """point_embed""",
    """pe_layer.positional_encoding_gaussian_matrix""": """shared_embedding.positional_embedding""",
    """image_encoder""": """vision_encoder""",
    """neck.0""": """neck.conv1""",
    """neck.1""": """neck.layer_norm1""",
    """neck.2""": """neck.conv2""",
    """neck.3""": """neck.layer_norm2""",
    """patch_embed.proj""": """patch_embed.projection""",
    """.norm""": """.layer_norm""",
    """blocks""": """layers""",
}






def 	UpperCamelCase	(							snake_case__  : Union[str, Any]	):




     '''simple docstring'''






     __snake_case      :Dict 							=    {}
     state_dict.pop("""pixel_mean"""						,snake_case__	)
     state_dict.pop("""pixel_std"""						,snake_case__	)

     __snake_case      :str 							=    R""".*.output_hypernetworks_mlps.(\d+).layers.(\d+).*"""

     for key, value in state_dict.items():
          for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
               if key_to_modify in key:
                    __snake_case      :Dict 							=    key.replace(snake_case__						,snake_case__	)

          if re.match(snake_case__						,snake_case__	):
               __snake_case      :Optional[Any] 							=    int(re.match(snake_case__						,snake_case__	).group(2	)	)
               if layer_nb == 0:
                    __snake_case      :List[Any] 							=    key.replace("""layers.0"""						,"""proj_in"""	)
               elif layer_nb == 1:
                    __snake_case      :List[str] 							=    key.replace("""layers.1"""						,"""layers.0"""	)
               elif layer_nb == 2:
                    __snake_case      :Optional[Any] 							=    key.replace("""layers.2"""						,"""proj_out"""	)

          __snake_case      :Union[str, Any] 							=    value

     __snake_case      :int 							=    model_state_dict[
         """prompt_encoder.shared_embedding.positional_embedding"""
     ]

     return model_state_dict






def 	UpperCamelCase	(							snake_case__  : Dict						,snake_case__  : Dict						,snake_case__  : List[str]						,snake_case__  : Dict="ybelkada/segment-anything"	):




     '''simple docstring'''






     __snake_case      :Union[str, Any] 							=    hf_hub_download(snake_case__						,f'''checkpoints/{model_name}.pth'''	)

     if "sam_vit_b" in model_name:
          __snake_case      :List[str] 							=    SamConfig()
     elif "sam_vit_l" in model_name:
          __snake_case      :Dict 							=    SamVisionConfig(
              hidden_size=1024						,num_hidden_layers=24						,num_attention_heads=16						,global_attn_indexes=[5, 11, 17, 23]						,)

          __snake_case      :Union[str, Any] 							=    SamConfig(
              vision_config=snake_case__						,)
     elif "sam_vit_h" in model_name:
          __snake_case      :List[str] 							=    SamVisionConfig(
              hidden_size=1280						,num_hidden_layers=32						,num_attention_heads=16						,global_attn_indexes=[7, 15, 23, 31]						,)

          __snake_case      :Optional[Any] 							=    SamConfig(
              vision_config=snake_case__						,)

     __snake_case      :Tuple 							=    torch.load(snake_case__						,map_location="""cpu"""	)
     __snake_case      :List[str] 							=    replace_keys(snake_case__	)

     __snake_case      :Optional[int] 							=    SamImageProcessor()

     __snake_case      :Dict 							=    SamProcessor(image_processor=snake_case__	)
     __snake_case      :Tuple 							=    SamModel(snake_case__	)

     hf_model.load_state_dict(snake_case__	)
     __snake_case      :Any 							=    hf_model.to("""cuda"""	)

     __snake_case      :Dict 							=    """https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"""
     __snake_case      :Optional[int] 							=    Image.open(requests.get(snake_case__						,stream=snake_case__	).raw	).convert("""RGB"""	)

     __snake_case      :str 							=    [[[400, 650]]]
     __snake_case      :str 							=    [[1]]

     __snake_case      :Tuple 							=    processor(images=np.array(snake_case__	)						,return_tensors="""pt"""	).to("""cuda"""	)

     with torch.no_grad():
          __snake_case      :Dict 							=    hf_model(**snake_case__	)
     __snake_case      :Optional[Any] 							=    output.iou_scores.squeeze()

     if model_name == "sam_vit_h_4b8939":
          assert scores[-1].item() == 0.5_7_9_8_9_0_2_5_1_1_5_9_6_6_8

          __snake_case      :Dict 							=    processor(
              images=np.array(snake_case__	)						,input_points=snake_case__						,input_labels=snake_case__						,return_tensors="""pt"""	).to("""cuda"""	)

          with torch.no_grad():
               __snake_case      :List[Any] 							=    hf_model(**snake_case__	)
          __snake_case      :List[Any] 							=    output.iou_scores.squeeze()

          assert scores[-1].item() == 0.9_7_1_2_6_0_3_0_9_2_1_9_3_6_0_4

          __snake_case      :Any 							=    ((75, 275, 1725, 850),)

          __snake_case      :Tuple 							=    processor(images=np.array(snake_case__	)						,input_boxes=snake_case__						,return_tensors="""pt"""	).to("""cuda"""	)

          with torch.no_grad():
               __snake_case      :int 							=    hf_model(**snake_case__	)
          __snake_case      :Any 							=    output.iou_scores.squeeze()

          assert scores[-1].item() == 0.8_6_8_6_0_1_5_6_0_5_9_2_6_5_1_4

          # Test with 2 points and 1 image.
          __snake_case      :Optional[int] 							=    [[[400, 650], [800, 650]]]
          __snake_case      :Any 							=    [[1, 1]]

          __snake_case      :List[str] 							=    processor(
              images=np.array(snake_case__	)						,input_points=snake_case__						,input_labels=snake_case__						,return_tensors="""pt"""	).to("""cuda"""	)

          with torch.no_grad():
               __snake_case      :int 							=    hf_model(**snake_case__	)
          __snake_case      :Dict 							=    output.iou_scores.squeeze()

          assert scores[-1].item() == 0.9_9_3_6_0_4_7_7_9_2_4_3_4_6_9_2


if __name__ == "__main__":
     lowerCamelCase__       		=  argparse.ArgumentParser()
     lowerCamelCase__       		=  ["""sam_vit_b_01ec64""", """sam_vit_h_4b8939""", """sam_vit_l_0b3195"""]
     parser.add_argument(
         """--model_name""",
         default="""sam_vit_h_4b8939""",
         choices=choices,
         type=str,
         help="""Path to hf config.json of model to convert""",
     )
     parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
     parser.add_argument(
         """--push_to_hub""",
         action="""store_true""",
         help="""Whether to push the model and processor to the hub after converting""",
     )
     parser.add_argument(
         """--model_hub_id""",
         default="""ybelkada/segment-anything""",
         choices=choices,
         type=str,
         help="""Path to hf config.json of model to convert""",
     )

     lowerCamelCase__       		=  parser.parse_args()

     convert_sam_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.model_hub_id)










from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)






def 	UpperCamelCase	(							snake_case__  : List[Any]						,snake_case__  : List[str]	):




     '''simple docstring'''






     __snake_case      :int 							=    []
     for part_id in partition_order:
          __snake_case      :int 							=    df.where(f'''SPARK_PARTITION_ID() = {part_id}'''	).collect()
          for row_idx, row in enumerate(snake_case__	):
               expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())	)
     return expected_row_ids_and_row_dicts






@require_not_windows
@require_dill_gt_0_3_2
def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :List[str] 							=    pyspark.sql.SparkSession.builder.master("""local[*]"""	).appName("""pyspark"""	).getOrCreate()
     __snake_case      :Any 							=    spark.range(100	).repartition(1	)
     __snake_case      :int 							=    Spark(snake_case__	)
     # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
     # that each partition can hold 2 rows.
     spark_builder._repartition_df_if_needed(max_shard_size=16	)
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 50






@require_not_windows
@require_dill_gt_0_3_2
def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Tuple 							=    pyspark.sql.SparkSession.builder.master("""local[*]"""	).appName("""pyspark"""	).getOrCreate()
     __snake_case      :int 							=    spark.range(10	).repartition(2	)
     __snake_case      :str 							=    [1, 0]
     __snake_case      :List[Any] 							=    _generate_iterable_examples(snake_case__						,snake_case__	)  # Reverse the partitions.
     __snake_case      :Union[str, Any] 							=    _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__						,snake_case__	)

     for i, (row_id, row_dict) in enumerate(generate_fn()	):
          __snake_case							,			__snake_case      :Union[str, Any] 							=    expected_row_ids_and_row_dicts[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Union[str, Any] 							=    pyspark.sql.SparkSession.builder.master("""local[*]"""	).appName("""pyspark"""	).getOrCreate()
     __snake_case      :Tuple 							=    spark.range(10	).repartition(1	)
     __snake_case      :Dict 							=    SparkExamplesIterable(snake_case__	)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(snake_case__	):
          assert row_id == f'''0_{i}'''
          assert row_dict == {"id": i}






@require_not_windows
@require_dill_gt_0_3_2
def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :List[str] 							=    pyspark.sql.SparkSession.builder.master("""local[*]"""	).appName("""pyspark"""	).getOrCreate()
     __snake_case      :Union[str, Any] 							=    spark.range(30	).repartition(3	)
     # Mock the generator so that shuffle reverses the partition indices.
     with patch("""numpy.random.Generator"""	) as generator_mock:
          __snake_case      :Dict 							=    lambda snake_case__	: x.reverse()
          __snake_case      :int 							=    _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__						,[2, 1, 0]	)

          __snake_case      :Dict 							=    SparkExamplesIterable(snake_case__	).shuffle_data_sources(snake_case__	)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(snake_case__	):
               __snake_case							,			__snake_case      :List[str] 							=    expected_row_ids_and_row_dicts[i]
               assert row_id == expected_row_id
               assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Union[str, Any] 							=    pyspark.sql.SparkSession.builder.master("""local[*]"""	).appName("""pyspark"""	).getOrCreate()
     __snake_case      :Tuple 							=    spark.range(20	).repartition(4	)

     # Partitions 0 and 2
     __snake_case      :List[Any] 							=    SparkExamplesIterable(snake_case__	).shard_data_sources(worker_id=0						,num_workers=2	)
     assert shard_it_a.n_shards == 2
     __snake_case      :Dict 							=    _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__						,[0, 2]	)
     for i, (row_id, row_dict) in enumerate(snake_case__	):
          __snake_case							,			__snake_case      :Tuple 							=    expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     __snake_case      :str 							=    SparkExamplesIterable(snake_case__	).shard_data_sources(worker_id=1						,num_workers=2	)
     assert shard_it_a.n_shards == 2
     __snake_case      :Optional[int] 							=    _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__						,[1, 3]	)
     for i, (row_id, row_dict) in enumerate(snake_case__	):
          __snake_case							,			__snake_case      :Dict 							=    expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Union[str, Any] 							=    pyspark.sql.SparkSession.builder.master("""local[*]"""	).appName("""pyspark"""	).getOrCreate()
     __snake_case      :Tuple 							=    spark.range(100	).repartition(1	)
     __snake_case      :Dict 							=    Spark(snake_case__	)
     # Choose a small max_shard_size for maximum partitioning.
     spark_builder._repartition_df_if_needed(max_shard_size=1	)
     # The new number of partitions should not be greater than the number of rows.
     assert spark_builder.df.rdd.getNumPartitions() == 100


