






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	(						)      ->					int:
    _UpperCAmelCase				:      Optional[Any]			  =				'''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
    _UpperCAmelCase				:      Dict			  =				Image.open(requests.get(__A		,		stream=__A ).raw ).convert('''RGB''' )

    return image

def      _lowerCamelCase	(						__A    :							Optional[int] )      ->					int:
    _UpperCAmelCase				:      Optional[int]			  =				[]
    # 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    :							Tuple		,		__A    :							List[Any]		,		__A    :							Optional[int] )      ->					List[Any]:
    _UpperCAmelCase				:      List[str]			  =				dct.pop(__A )
    _UpperCAmelCase				:      str			  =				val

def      _lowerCamelCase	(						__A    :							Optional[int]		,		__A    :							Tuple )      ->					Dict:
    for i in range(config.vision_config.num_hidden_layers ):
        # read in original q and v biases
        _UpperCAmelCase				:      Optional[int]			  =				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias''' )
        _UpperCAmelCase				:      Dict			  =				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias''' )

        # next, set bias in the state dict
        _UpperCAmelCase				:      Optional[Any]			  =				torch.cat((q_bias, torch.zeros_like(__A		,		requires_grad=__A ), v_bias) )
        _UpperCAmelCase				:      str			  =				qkv_bias

def      _lowerCamelCase	(						__A    :							str )      ->					str:
    _UpperCAmelCase				:      Union[str, Any]			  =				364 if '''coco''' in model_name else 224
    _UpperCAmelCase				:      Tuple			  =				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				:      Union[str, Any]			  =				TaConfig.from_pretrained('''google/flan-t5-xl'''		,		dense_act_fn='''gelu'''		,		bos_token_id=1 ).to_dict()
    elif "t5-xxl" in model_name:
        _UpperCAmelCase				:      Dict			  =				TaConfig.from_pretrained('''google/flan-t5-xxl'''		,		dense_act_fn='''gelu'''		,		bos_token_id=1 ).to_dict()
    elif "vicuna-7b" in model_name:
        _UpperCAmelCase				:      int			  =				LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''		,		vocab_size=32_001 ).to_dict()
    elif "vicuna-13b" in model_name:
        _UpperCAmelCase				:      Any			  =				LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''		,		vocab_size=32_001 ).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				:      Optional[Any]			  =				InstructBlipQFormerConfig(vocab_size=30_523 ).to_dict()
    _UpperCAmelCase				:      Tuple			  =				InstructBlipConfig(vision_config=__A		,		text_config=__A		,		qformer_config=__A )

    return config, image_size

@torch.no_grad()
def      _lowerCamelCase	(						__A    :							List[str]		,		__A    :							int=None		,		__A    :							List[str]=False )      ->					str:
    _UpperCAmelCase				:      str			  =				AutoTokenizer.from_pretrained('''bert-base-uncased'''		,		truncation_side='''left''' )
    qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''} )

    if "t5" in model_name:
        _UpperCAmelCase				:      Union[str, Any]			  =				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				:      Optional[Any]			  =				LlamaTokenizerFast.from_pretrained(
            '''huggyllama/llama-7b'''		,		truncation_side='''left'''		,		bos_token='''</s>'''		,		unk_token='''</s>''' )
        tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''} )

    _UpperCAmelCase  , _UpperCAmelCase				:      List[str]			  =				get_blipa_config(__A )
    _UpperCAmelCase				:      Tuple			  =				InstructBlipForConditionalGeneration(__A ).eval()

    _UpperCAmelCase				:      List[Any]			  =				{
        '''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				:      Tuple			  =				model_name_to_original[model_name]

    # load original model
    print('''Loading original model...''' )
    _UpperCAmelCase				:      int			  =				'''cuda:1''' if torch.cuda.is_available() else '''cpu'''
    _UpperCAmelCase				:      List[str]			  =				'''cuda:2''' if torch.cuda.is_available() else '''cpu'''
    _UpperCAmelCase  , _UpperCAmelCase  , _UpperCAmelCase				:      str			  =				load_model_and_preprocess(
        name=__A		,		model_type=__A		,		is_eval=__A		,		device=__A )
    original_model.eval()
    print('''Done!''' )

    # update state dict keys
    _UpperCAmelCase				:      Dict			  =				original_model.state_dict()
    _UpperCAmelCase				:      Optional[int]			  =				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				:      Optional[Any]			  =				state_dict.pop(__A )
        if key.startswith('''Qformer.bert''' ):
            _UpperCAmelCase				:      Any			  =				key.replace('''Qformer.bert'''		,		'''qformer''' )
        if "attention.self" in key:
            _UpperCAmelCase				:      List[Any]			  =				key.replace('''self'''		,		'''attention''' )
        if "llm_proj" in key:
            _UpperCAmelCase				:      Optional[Any]			  =				key.replace('''llm_proj'''		,		'''language_projection''' )
        if "t5_proj" in key:
            _UpperCAmelCase				:      Tuple			  =				key.replace('''t5_proj'''		,		'''language_projection''' )
        if key.startswith('''llm_model''' ):
            _UpperCAmelCase				:      Optional[Any]			  =				key.replace('''llm_model'''		,		'''language_model''' )
        if key.startswith('''t5''' ):
            _UpperCAmelCase				:      Optional[int]			  =				key.replace('''t5'''		,		'''language''' )
        _UpperCAmelCase				:      Tuple			  =				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				:      str			  =				load_demo_image()
    _UpperCAmelCase				:      Union[str, Any]			  =				'''What is unusual about this image?'''

    # create processor
    _UpperCAmelCase				:      List[str]			  =				BlipImageProcessor(
        size={'''height''': image_size, '''width''': image_size}		,		image_mean=__A		,		image_std=__A )
    _UpperCAmelCase				:      str			  =				InstructBlipProcessor(
        image_processor=__A		,		tokenizer=__A		,		qformer_tokenizer=__A		,		)
    _UpperCAmelCase				:      Dict			  =				processor(images=__A		,		text=__A		,		return_tensors='''pt''' ).to(__A )

    # make sure processor creates exact same pixel values
    _UpperCAmelCase				:      List[str]			  =				vis_processors['''eval'''](__A ).unsqueeze(0 ).to(__A )
    _UpperCAmelCase				:      Union[str, Any]			  =				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				:      str			  =				original_model({'''image''': original_pixel_values, '''text_input''': [prompt]} ).logits
            _UpperCAmelCase				:      Any			  =				hf_model(**__A ).logits
        else:
            _UpperCAmelCase				:      str			  =				original_model(
                {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']} ).logits
            _UpperCAmelCase				:      Any			  =				tokenizer('''\n'''		,		return_tensors='''pt''' ).input_ids.to(__A )
            _UpperCAmelCase				:      Tuple			  =				label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id		,		-100 )
            _UpperCAmelCase				:      str			  =				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				:      Optional[Any]			  =				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				:      Any			  =				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				:      Optional[Any]			  =				hf_model.generate(
        **__A		,		do_sample=__A		,		num_beams=5		,		max_length=256		,		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				:      Optional[Any]			  =				2
    print('''Original generation:'''		,		__A )
    _UpperCAmelCase				:      str			  =				processor.batch_decode(__A		,		skip_special_tokens=__A )
    _UpperCAmelCase				:      Tuple			  =				[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__":
  SCREAMING_SNAKE_CASE 							=						argparse.ArgumentParser()
  SCREAMING_SNAKE_CASE 							=						[
      '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',
  )

  SCREAMING_SNAKE_CASE 							=						parser.parse_args()

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












# Logistic Regression from scratch

# In[62]:

# In[63]:

# importing all the required libraries

import numpy as np
from matplotlib import pyplot as plt
from sklearn import datasets

def      _lowerCamelCase	(						__A    :							Optional[int] )      ->					str:
    return 1 / (1 + np.exp(-z ))

def      _lowerCamelCase	(						__A    :							Tuple		,		__A    :							List[Any] )      ->					Optional[int]:
    return (-y * np.log(__A ) - (1 - y) * np.log(1 - h )).mean()

def      _lowerCamelCase	(						__A    :							Tuple		,		__A    :							List[Any]		,		__A    :							int )      ->					Any:
    _UpperCAmelCase				:      int			  =				np.dot(__A		,		__A )
    return np.sum(y * scores - np.log(1 + np.exp(__A ) ) )

def      _lowerCamelCase	(						__A    :							List[Any]		,		__A    :							Union[str, Any]		,		__A    :							str		,		__A    :							Dict=70_000 )      ->					Optional[Any]:
    _UpperCAmelCase				:      Union[str, Any]			  =				np.zeros(x.shape[1] )

    for iterations in range(__A ):
        _UpperCAmelCase				:      List[Any]			  =				np.dot(__A		,		__A )
        _UpperCAmelCase				:      Union[str, Any]			  =				sigmoid_function(__A )
        _UpperCAmelCase				:      Optional[Any]			  =				np.dot(x.T		,		h - y ) / y.size
        _UpperCAmelCase				:      Optional[Any]			  =				theta - alpha * gradient  # updating the weights
        _UpperCAmelCase				:      List[Any]			  =				np.dot(__A		,		__A )
        _UpperCAmelCase				:      Any			  =				sigmoid_function(__A )
        _UpperCAmelCase				:      Union[str, Any]			  =				cost_function(__A		,		__A )
        if iterations % 100 == 0:
            print(f'''loss: {j} \t''' )  # printing the loss after every 100 iterations
    return theta


# In[68]:

if __name__ == "__main__":
  SCREAMING_SNAKE_CASE 							=						datasets.load_iris()
  SCREAMING_SNAKE_CASE 							=						iris.data[:, :2]
  SCREAMING_SNAKE_CASE 							=						(iris.target != 0) * 1

  SCREAMING_SNAKE_CASE 							=						0.1
  SCREAMING_SNAKE_CASE 							=						logistic_reg(alpha, x, y, max_iterations=70000)
  print('theta: ', theta)  # printing the theta i.e our weights vector

  def      _lowerCamelCase	(						__A    :							Any )      ->					int:
      return sigmoid_function(
          np.dot(__A		,		__A ) )  # predicting the value of probability from the logistic regression algorithm

  plt.figure(figsize=(10, 6))
  plt.scatter(x[y == 0][:, 0], x[y == 0][:, 1], color='b', label='0')
  plt.scatter(x[y == 1][:, 0], x[y == 1][:, 1], color='r', label='1')
  ((SCREAMING_SNAKE_CASE)							,			(SCREAMING_SNAKE_CASE)) 							=						(x[:, 0].min(), x[:, 0].max())
  ((SCREAMING_SNAKE_CASE)							,			(SCREAMING_SNAKE_CASE)) 							=						(x[:, 1].min(), x[:, 1].max())
  ((SCREAMING_SNAKE_CASE)							,			(SCREAMING_SNAKE_CASE)) 							=						np.meshgrid(np.linspace(xa_min, xa_max), np.linspace(xa_min, xa_max))
  SCREAMING_SNAKE_CASE 							=						np.c_[xxa.ravel(), xxa.ravel()]
  SCREAMING_SNAKE_CASE 							=						predict_prob(grid).reshape(xxa.shape)
  plt.contour(xxa, xxa, probs, [0.5], linewidths=1, colors='black')

  plt.legend()


  plt.show()




