def __init__(self, model_name, auth_token=None):

self.device = "cuda" if torch.cuda.is_available() else "cpu"

self.model_name = model_name

self.model, self.preprocess, self.model_hash = self._load_model(model_name, auth_token=auth_token)

self.model = self.model.to(self.device)

def _load_model(self,

name: str,

device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",

auth_token=None):

model = CLIPModel.from_pretrained(name, use_auth_token=auth_token)

preprocessing = CLIPProcessor.from_pretrained(name, use_auth_token=auth_token)

return model, preprocessing, hash

def encode_images(self, images: Union[List[str], List[PIL.Image.Image]], batch_size: int):

def transform_fn(el):

imgs = el['image'] if isinstance(el['image'][0], PIL.Image.Image) else [Image().decode_example(_) for _ in

el['image']]

return self.preprocess(images=imgs, return_tensors='pt')

dataset = Dataset.from_dict({'image': images})

dataset = dataset.cast_column('image', Image(decode=False)) if isinstance(images[0], str) else dataset

# dataset = dataset.map(map_fn,

# batched=True,

# remove_columns=['image'])

dataset.set_format('torch')

dataset.set_transform(transform_fn)

dataloader = DataLoader(dataset, batch_size=batch_size)

image_embeddings = []

pbar = tqdm(total=len(images) // batch_size, position=0)

with torch.no_grad():

for batch in dataloader:

batch = {k: v.to(self.device) for k, v in batch.items()}

image_embeddings.extend(self.model.get_image_features(**batch).detach().cpu().numpy())

pbar.update(1)

pbar.close()

return np.stack(image_embeddings)

def encode_text(self, text: List[str], batch_size: int):

dataset = Dataset.from_dict({'text': text})

dataset = dataset.map(lambda el: self.preprocess(text=el['text'], return_tensors="pt",

max_length=77, padding="max_length", truncation=True),

batched=True,

remove_columns=['text'])

dataset.set_format('torch')

dataloader = DataLoader(dataset, batch_size=batch_size)

text_embeddings = []

pbar = tqdm(total=len(text) // batch_size, position=0)

with torch.no_grad():

for batch in dataloader:

batch = {k: v.to(self.device) for k, v in batch.items()}

text_embeddings.extend(self.model.get_text_features(**batch).detach().cpu().numpy())

pbar.update(1)

pbar.close()

return np.stack(text_embeddings)

def _cosine_similarity(self, key_vectors: np.ndarray, space_vectors: np.ndarray, normalize=True):

if normalize:

key_vectors = key_vectors / np.linalg.norm(key_vectors, ord=2, axis=-1, keepdims=True)

return np.matmul(key_vectors, space_vectors.T)

def _nearest_neighbours(self, k, key_vectors, space_vectors, normalize=True, debug=False):

if type(key_vectors) == List:

key_vectors = np.array(key_vectors)

if type(space_vectors) == List:

space_vectors = np.array(space_vectors)

cosine_sim = self._cosine_similarity(key_vectors, space_vectors, normalize=normalize)

nn = cosine_sim.argsort()[:, -k:][:, ::-1]

return nn

def zero_shot_classification(self, images, text_labels: List[str], debug=False):

"""

# encode text

text_vectors = self.encode_text(text_labels, batch_size=8)

# encode images

image_vectors = self.encode_images(images, batch_size=8)

# compute cosine similarity

cosine_sim = self._cosine_similarity(image_vectors, text_vectors)

if debug:

print(cosine_sim)

preds = np.argmax(cosine_sim, axis=-1)

return [text_labels[idx] for idx in preds]

def retrieval(self, queries: List[str], top_k: int = 10):

"""

# encode text

text_vectors = self.encode_text(queries, batch_size=8)

# compute cosine similarity

# cosine_sim = self._cosine_similarity(text_vectors, self.image_vectors)

return self._nearest_neighbours(k=top_k, key_vectors=text_vectors, space_vectors=self.image_vectors)

# return np.argmax(cosine_sim, axis=-1)