from transformers import pipeline
from huggingface_hub import create_repo, upload_folder, snapshot_download
import torch
import transformers
from transformers import AutoModelForCausalLM
import os
from pathlib import Path

model_id = 'Qwen/Qwen-VL-Chat'
new_repo_id = 'yujiepan/qwen-vl-tiny-random'


def replace_in_file(file_path, old: str, new: str):
    with open(file_path, 'r', encoding='utf-8') as f:
        visual_code = f.read()
    visual_code = visual_code.replace(old, new)
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(visual_code)


def download_modeling_codes():
    snapshot_download(repo_id=model_id, allow_patterns='*.py',
                      local_dir='./qwen_vl_tiny_random', local_dir_use_symlinks=False)
    # The hard coded "128" is changed for smaller model size.
    replace_in_file('./qwen_vl_tiny_random/visual.py',
                    'num_heads=output_dim // 128,', 'num_heads=output_dim // 4,')


def create_config():
    from qwen_vl_tiny_random.configuration_qwen import QWenConfig
    config = QWenConfig()
    config.fp16 = True
    config.hidden_size = 8
    config.intermediate_size = 16
    config.kv_channels = 4
    config.num_attention_heads = 2
    config.num_hidden_layers = 2
    config.seq_length = 2048

    config.visual  = {
        "heads": 2,
        "image_size": 448,
        "image_start_id": 151857,
        "layers": 2,
        "mlp_ratio": 1.0,
        "output_dim": 8,
        "patch_size": 14,
        "width": 8,
    }
    print(config)
    return config


def create_model(config):
    from qwen_vl_tiny_random.modeling_qwen import QWenLMHeadModel, QWenModel
    from qwen_vl_tiny_random.configuration_qwen import QWenConfig
    from transformers import AutoModelForCausalLM, AutoConfig, AutoModel
    AutoConfig.register("qwen", QWenConfig)
    AutoModel.register(QWenConfig, QWenModel)
    AutoModelForCausalLM.register(QWenConfig, QWenLMHeadModel)
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
    model.generation_config = transformers.GenerationConfig.from_pretrained(
        model_id, trust_remote_code=True)
    return model


def try_inference(model, tokenizer):
    model = model.cuda()
    query = tokenizer.from_list_format([
        {'image': 'https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg'},
        {'text': '这是什么'},
    ])
    response, history = model.chat(tokenizer, query=query, history=None)
    print(response)


download_modeling_codes()
config = create_config()
model = create_model(config)
tokenizer = transformers.AutoTokenizer.from_pretrained(
    model_id, trust_remote_code=True)
try_inference(model, tokenizer)

model.save_pretrained('./qwen_vl_tiny_random/')
tokenizer.save_pretrained('./qwen_vl_tiny_random/')

create_repo(new_repo_id, exist_ok=True)
upload_folder(repo_id=new_repo_id, folder_path='./qwen_vl_tiny_random/',
              ignore_patterns='__pycache__')

model = transformers.AutoModelForCausalLM.from_pretrained(
    new_repo_id, trust_remote_code=True).cuda()
tokenizer = transformers.AutoTokenizer.from_pretrained(
    new_repo_id, trust_remote_code=True)
try_inference(model, tokenizer)