'''simple docstring'''





from __future__ import annotations

from functools import lru_cache
from math import ceil

a_							=      100

a_							=      set(range(3, NUM_PRIMES, 2))
primes.add(2)
a_							=      42

for prime in range(3, ceil(NUM_PRIMES**0.5), 2):
       if prime not in primes:
              continue
       primes.difference_update(set(range(prime * prime, NUM_PRIMES, prime)))




@lru_cache(maxsize=1_00					)
def       __lowerCAmelCase      (						A_   : int					)    ->  Dict:
  if number_to_partition < 0:
    return set()
  elif number_to_partition == 0:
    return {1}

  __UpperCAmelCase						    =       set()
  __UpperCAmelCase						    =       42
  __UpperCAmelCase						    =       42

  for prime in primes:
    if prime > number_to_partition:
      continue
    for sub in partition(number_to_partition - prime					):
      ret.add(sub * prime					)

  return ret




def       __lowerCAmelCase      (						A_   : int = 50_00					)    ->  Union[str, Any]:
  for number_to_partition in range(1						,	lowerCamelCase_					):
    if len(partition(lowerCamelCase_					)					) > number_unique_partitions:
      return number_to_partition
  return None


if __name__ == "__main__":
       print(F"{solution() = }")




import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD




def       __lowerCAmelCase      (						)    ->  Optional[Any]:
  __UpperCAmelCase						    =       "https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"
  __UpperCAmelCase						    =       Image.open(requests.get(A_						,	stream=A_					).raw					).convert("RGB"					)

  return image




def       __lowerCAmelCase      (						A_   : List[Any]					)    ->  List[str]:
  __UpperCAmelCase						    =       []
  # fmt: off

  # vision encoder
  rename_keys.append(("visual_encoder.cls_token", "vision_model.embeddings.class_embedding")					)
  rename_keys.append(("visual_encoder.pos_embed", "vision_model.embeddings.position_embedding")					)
  rename_keys.append(("visual_encoder.patch_embed.proj.weight", "vision_model.embeddings.patch_embedding.weight")					)
  rename_keys.append(("visual_encoder.patch_embed.proj.bias", "vision_model.embeddings.patch_embedding.bias")					)
  rename_keys.append(("ln_vision.weight", "vision_model.post_layernorm.weight")					)
  rename_keys.append(("ln_vision.bias", "vision_model.post_layernorm.bias")					)

  for i in range(config.vision_config.num_hidden_layers					):
    rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.weight''', F'''vision_model.encoder.layers.{i}.layer_norm1.weight''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.bias''', F'''vision_model.encoder.layers.{i}.layer_norm1.bias''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.weight''', F'''vision_model.encoder.layers.{i}.layer_norm2.weight''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.bias''', F'''vision_model.encoder.layers.{i}.layer_norm2.bias''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.attn.qkv.weight''', F'''vision_model.encoder.layers.{i}.self_attn.qkv.weight''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.weight''', F'''vision_model.encoder.layers.{i}.self_attn.projection.weight''',)					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.bias''', F'''vision_model.encoder.layers.{i}.self_attn.projection.bias''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc1.weight''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc1.bias''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc2.weight''')					)
    rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc2.bias''')					)

  # QFormer
  rename_keys.append(("Qformer.bert.embeddings.LayerNorm.weight", "qformer.embeddings.layernorm.weight")					)
  rename_keys.append(("Qformer.bert.embeddings.LayerNorm.bias", "qformer.embeddings.layernorm.bias")					)

  # fmt: on
  return rename_keys




def       __lowerCAmelCase      (						A_   : Optional[int]						,	A_   : int						,	A_   : str					)    ->  List[Any]:
  __UpperCAmelCase						    =       dct.pop(A_					)
  __UpperCAmelCase						    =       val




def       __lowerCAmelCase      (						A_   : Optional[int]						,	A_   : Optional[int]					)    ->  int:
  for i in range(config.vision_config.num_hidden_layers					):
    # read in original q and v biases
    __UpperCAmelCase						    =       state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias'''					)
    __UpperCAmelCase						    =       state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''					)

    # next, set bias in the state dict
    __UpperCAmelCase						    =       torch.cat((q_bias, torch.zeros_like(A_						,	requires_grad=A_					), v_bias)					)
    __UpperCAmelCase						    =       qkv_bias




def       __lowerCAmelCase      (						A_   : Any					)    ->  int:
  __UpperCAmelCase						    =       3_64 if "coco" in model_name else 2_24
  __UpperCAmelCase						    =       InstructBlipVisionConfig(image_size=A_					).to_dict()

  # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
  # seems like flan-T5 models don't have bos_token_id properly set?
  if "t5-xl" in model_name:
    __UpperCAmelCase						    =       TaConfig.from_pretrained("google/flan-t5-xl"						,	dense_act_fn="gelu"						,	bos_token_id=1					).to_dict()
  elif "t5-xxl" in model_name:
    __UpperCAmelCase						    =       TaConfig.from_pretrained("google/flan-t5-xxl"						,	dense_act_fn="gelu"						,	bos_token_id=1					).to_dict()
  elif "vicuna-7b" in model_name:
    __UpperCAmelCase						    =       LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf"						,	vocab_size=3_20_01					).to_dict()
  elif "vicuna-13b" in model_name:
    __UpperCAmelCase						    =       LlamaConfig.from_pretrained("decapoda-research/llama-13b-hf"						,	vocab_size=3_20_01					).to_dict()
  else:
    raise ValueError("Model name not supported"					)

  # the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
  __UpperCAmelCase						    =       InstructBlipQFormerConfig(vocab_size=3_05_23					).to_dict()
  __UpperCAmelCase						    =       InstructBlipConfig(vision_config=A_						,	text_config=A_						,	qformer_config=A_					)

  return config, image_size




@torch.no_grad()
def       __lowerCAmelCase      (						A_   : int						,	A_   : Union[str, Any]=None						,	A_   : Optional[Any]=False					)    ->  Dict:
  __UpperCAmelCase						    =       AutoTokenizer.from_pretrained("bert-base-uncased"						,	truncation_side="left"					)
  qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"}					)

  if "t5" in model_name:
    __UpperCAmelCase						    =       TaTokenizerFast.from_pretrained("google/flan-t5-xl"						,	truncation_side="left"					)
  elif "vicuna" in model_name:
    # the following was used in the original implementation:
    # tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
    # tokenizer.add_special_tokens({"pad_token": "[PAD]"})
    # tokenizer.add_special_tokens({"bos_token": "</s>"})
    # tokenizer.add_special_tokens({"eos_token": "</s>"})
    # tokenizer.add_special_tokens({"unk_token": "</s>"})
    __UpperCAmelCase						    =       LlamaTokenizerFast.from_pretrained(
        "huggyllama/llama-7b"						,	truncation_side="left"						,	bos_token="</s>"						,	unk_token="</s>"					)
    tokenizer.add_special_tokens({"pad_token": "[PAD]"}					)

  __UpperCAmelCase    ,		__UpperCAmelCase						    =       get_blipa_config(A_					)
  __UpperCAmelCase						    =       InstructBlipForConditionalGeneration(A_					).eval()

  __UpperCAmelCase						    =       {
      "instructblip-vicuna-7b": ("blip2_vicuna_instruct", "vicuna7b"),
      "instructblip-vicuna-13b": ("blip2_vicuna_instruct", "vicuna13b"),
      "instructblip-flan-t5-xl": ("blip2_t5_instruct", "flant5xl"),
      "instructblip-flan-t5-xxl": ("blip2_t5_instruct", "flant5xxl"),
  }

  __UpperCAmelCase    ,		__UpperCAmelCase						    =       model_name_to_original[model_name]

  # load original model
  print("Loading original model..."					)
  __UpperCAmelCase						    =       "cuda:1" if torch.cuda.is_available() else "cpu"
  __UpperCAmelCase						    =       "cuda:2" if torch.cuda.is_available() else "cpu"
  __UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase						    =       load_model_and_preprocess(
      name=A_						,	model_type=A_						,	is_eval=A_						,	device=A_					)
  original_model.eval()
  print("Done!"					)

  # update state dict keys
  __UpperCAmelCase						    =       original_model.state_dict()
  __UpperCAmelCase						    =       create_rename_keys(A_					)
  for src, dest in rename_keys:
    rename_key(A_						,	A_						,	A_					)

  # some keys can be renamed efficiently
  for key, val in state_dict.copy().items():
    __UpperCAmelCase						    =       state_dict.pop(A_					)
    if key.startswith("Qformer.bert"					):
      __UpperCAmelCase						    =       key.replace("Qformer.bert"						,	"qformer"					)
    if "attention.self" in key:
      __UpperCAmelCase						    =       key.replace("self"						,	"attention"					)
    if "llm_proj" in key:
      __UpperCAmelCase						    =       key.replace("llm_proj"						,	"language_projection"					)
    if "t5_proj" in key:
      __UpperCAmelCase						    =       key.replace("t5_proj"						,	"language_projection"					)
    if key.startswith("llm_model"					):
      __UpperCAmelCase						    =       key.replace("llm_model"						,	"language_model"					)
    if key.startswith("t5"					):
      __UpperCAmelCase						    =       key.replace("t5"						,	"language"					)
    __UpperCAmelCase						    =       val

  # read in qv biases
  read_in_q_v_bias(A_						,	A_					)

  # note: weights get loaded in torch.float32 by default
  hf_model.load_state_dict(A_						,	strict=A_					)

  __UpperCAmelCase						    =       load_demo_image()
  __UpperCAmelCase						    =       "What is unusual about this image?"

  # create processor
  __UpperCAmelCase						    =       BlipImageProcessor(
      size={"height": image_size, "width": image_size}						,	image_mean=A_						,	image_std=A_					)
  __UpperCAmelCase						    =       InstructBlipProcessor(
      image_processor=A_						,	tokenizer=A_						,	qformer_tokenizer=A_						,	)
  __UpperCAmelCase						    =       processor(images=A_						,	text=A_						,	return_tensors="pt"					).to(A_					)

  # make sure processor creates exact same pixel values
  __UpperCAmelCase						    =       vis_processors["eval"](A_					).unsqueeze(0					).to(A_					)
  __UpperCAmelCase						    =       inputs.pixel_values
  assert torch.allclose(original_pixel_values.to(pixel_values.device					)						,	A_					)

  original_model.to(A_					)
  hf_model.to(A_					)
  with torch.no_grad():
    if "vicuna" in model_name:
      __UpperCAmelCase						    =       original_model({"image": original_pixel_values, "text_input": [prompt]}					).logits
      __UpperCAmelCase						    =       hf_model(**A_					).logits
    else:
      __UpperCAmelCase						    =       original_model(
          {"image": original_pixel_values, "text_input": [prompt], "text_output": ["\n"]}					).logits
      __UpperCAmelCase						    =       tokenizer("\n"						,	return_tensors="pt"					).input_ids.to(A_					)
      __UpperCAmelCase						    =       label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id						,	-1_00					)
      __UpperCAmelCase						    =       hf_model(**A_						,	labels=A_					).logits

  print("First values of original logits:"						,	original_logits[0, :3, :3]					)
  print("First values of HF logits:"						,	logits[0, :3, :3]					)

  # assert values
  assert original_logits.shape == logits.shape
  __UpperCAmelCase						    =       1e-4 if "vicuna" in model_name else 1e-5
  assert torch.allclose(original_logits.to(logits.device					)						,	A_						,	atol=A_					)
  print("Looks ok!"					)

  print("Generating with original model..."					)
  __UpperCAmelCase						    =       original_model.generate({"image": original_pixel_values, "prompt": prompt}						,	num_beams=5					)

  # important: we need to cast the weights of the HF model to the appropriate type
  print("Generating with HF model..."					)
  __UpperCAmelCase						    =       hf_model.generate(
      **A_						,	do_sample=A_						,	num_beams=5						,	max_length=2_56						,	min_length=1						,	top_p=0.9						,	repetition_penalty=1.5						,	length_penalty=1.0						,	temperature=1						,	)
  if "vicuna" in model_name:
    # convert output id 0 to 2 (eos_token_id)
    # TODO add this in the generate method?
    __UpperCAmelCase						    =       2
  print("Original generation:"						,	A_					)
  __UpperCAmelCase						    =       processor.batch_decode(A_						,	skip_special_tokens=A_					)
  __UpperCAmelCase						    =       [text.strip() for text in output_text]
  print("HF generation:"						,	A_					)

  if pytorch_dump_folder_path is not None:
    processor.save_pretrained(A_					)
    hf_model.save_pretrained(A_					)

  if push_to_hub:
    processor.push_to_hub(F'''Salesforce/{model_name}'''					)
    hf_model.push_to_hub(F'''Salesforce/{model_name}'''					)


if __name__ == "__main__":
       a_							=      argparse.ArgumentParser()
       a_							=      [
           """instructblip-vicuna-7b""",
           """instructblip-vicuna-13b""",
           """instructblip-flan-t5-xl""",
           """instructblip-flan-t5-xxl""",
       ]
       parser.add_argument(
           """--model_name""",
           default="""instructblip-flan-t5-xl""",
           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""",
       )

       a_							=      parser.parse_args()

       convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



