%%skip not $to_quantize.value 

import torch 
from datasets import load_dataset 
from tqdm.notebook import tqdm 
import requests 
from io import BytesIO 
import numpy as np 
from PIL import Image 
from requests.packages.urllib3.exceptions import InsecureRequestWarning 
requests.packages.urllib3.disable_warnings(InsecureRequestWarning) 

def check_text_data(data):
    """ 
    Check if the given data is text-based.
    """ 
    if isinstance(data, str): 
        return True 
    if isinstance(data, list): 
        return all(isinstance(x, str) for x in data) 
    return False 
}
def collate_fn_text(example, text_column="caption"):
     """ 
     Preprocesses an example by loading and transforming text data.
     Checks if the text data in the example is valid by calling the `check_text_data` function.
     If there is any error during the download process, returns None.
     Returns the preprocessed inputs with transformed image and text data.
     """ 
    assert len(example) == 1 
    example = example[0] 

    if not check_text_data(example[text_column]): 
        raise ValueError("Text data is not valid") 

    text_input = tokenizer( 
        example[text_column], 
        return_tensors='pt', 
        **tokenizer_kwargs) 

    return text_input 

def prepare_calibration_data_text(dataloader, init_steps):
     """ 
     This function prepares calibration data from a dataloader for a specified number of initialization steps.
     It iterates over the dataloader, fetching batches and storing the relevant data.
     """ 
    data = [] 
    print(f"Fetching {init_steps} samples for the initialization...") 
    with tqdm(total=init_steps) as pbar: 
        for batch in dataloader: 
            if len(data) == init_steps: 
                break 
            if batch: 
                pbar.update(1) 
                with torch.no_grad(): 
                    data.append(batch["input_ids"].to("cpu")) 
    return data

%%skip not $to_quantize.value 

import logging 
import nncf 

dataset = load_dataset("google-research-datasets/conceptual_captions", trust_remote_code=True) 
train_dataset = dataset["train"].shuffle(seed=42) 

dataloader_text = torch.utils.data.DataLoader(train_dataset, collate_fn=collate_fn_text, batch_size=1) 
calibration_data_text = prepare_calibration_data_text(dataloader_text, 50)

%%skip not $to_quantize.value 

def get_pil_from_url(url):
    """ 
    Downloads and converts an image from a URL to a PIL Image object.
    """ 
    response = requests.get(url, verify=False, timeout=20) 
    image = Image.open(BytesIO(response.content)) 
    return image.convert("RGB") 

def collate_fn_vision(example, image_column="image_url"):
    """ 
    Preprocesses an example by loading and transforming image data.
    Downloads the image specified by the URL in the image_column by calling the `get_pil_from_url` function.
    If there is any error during the download process, returns None.
    Returns the preprocessed inputs with transformed image and text data.
    """ 
    assert len(example) == 1 
    example = example[0] 

    url = example[image_column] 
    try: 
        image = get_pil_from_url(url) 
        h, w = image.size 
        if h == 1 or w == 1: 
            return None 
    except Exception: 
        return None 

    vision_input = processor(images=[image]) 
    return vision_input 

def prepare_calibration_data_vis(dataloader, init_steps):
    """ 
    This function prepares calibration data from a dataloader for a specified number of initialization steps.
    It iterates over the dataloader, fetching batches and storing the relevant data.
    """ 
    data = [] 
    print(f"Fetching {init_steps} samples for the initialization...") 
    with tqdm(total=init_steps) as pbar: 
        for batch in dataloader: 
            if len(data) == init_steps: 
                break 
            if batch: 
                pbar.update(1) 
                with torch.no_grad(): 
                    data.append(batch["pixel_values"].to("cpu")) 
    return data

%%skip not $to_quantize.value 

dataset = load_dataset("google-research-datasets/conceptual_captions", trust_remote_code=True) 
train_dataset = dataset["train"].shuffle(seed=42) 

dataloader_vis = torch.utils.data.DataLoader(train_dataset, collate_fn=collate_fn_vision, batch_size=1) 
calibration_data_vision = prepare_calibration_data_vis(dataloader_vis, 50)

int8_text_model_path = "jina-clip-text_v1_int8.xml"

%%skip not $to_quantize.value 

if len(calibration_data_text) == 0: 
    raise RuntimeError( 
        'Calibration dataset is empty.Please check internet connection and try to download images manually.'     ) 

ov_model_text = core.read_model(fp16_text_model_path) 

calibration_dataset = nncf.Dataset(calibration_data_text) 
quantized_model = nncf.quantize( 
    model=ov_model_text, 
    calibration_dataset=calibration_dataset 
) 
ov.save_model(quantized_model, int8_text_model_path)

int8_vision_model_path = "jina-clip-vision_v1_int8.xml"

%%skip not $to_quantize.value 

if len(calibration_data_vision) == 0: 
    raise RuntimeError( 
        'Calibration dataset is empty.Please check internet connection and try to download images manually.'     ) 

ov_model_vision = core.read_model(fp16_vision_model_path) 

calibration_dataset = nncf.Dataset(calibration_data_vision) 
quantized_model = nncf.quantize( 
    model=ov_model_vision, 
    calibration_dataset=calibration_dataset 
) 
ov.save_model(quantized_model, int8_vision_model_path)

%%skip not $to_quantize.value 

compiled_text_model_int8 = core.compile_model(int8_text_model_path, device.value) 
compiled_vision_model_int8 = core.compile_model(int8_vision_model_path, device.value) 

text_ov_res_int8 = compiled_text_model_int8(text_inputs["input_ids"]) 
vis_ov_res_int8 = compiled_vision_model_int8(vision_inputs["pixel_values"]) 

res = calc_simularity_softmax(vis_ov_res_int8[0], text_ov_res_int8[0]) 
visionize_result(img_furseal, TEXT_INPUTS, np.array(res[0]))