Spaces:

osanchik
/

PicFinder

Runtime error

modified faiss

d738c76 over 3 years ago

3.65 kB

from transformers import CLIPModel, CLIPTokenizer
from sklearn.metrics.pairwise import cosine_similarity
import faiss

from dataframe  import *

def get_model_info(model_ID, device):
    # Save the model to device
	model = CLIPModel.from_pretrained(model_ID).to(device)

    # Get the tokenizer
	tokenizer = CLIPTokenizer.from_pretrained(model_ID)

    # Return model, processor & tokenizer
	return model, tokenizer


def get_single_text_embedding(text, model, tokenizer, device):
    inputs = tokenizer(text, return_tensors = "pt", max_length=77, truncation=True).to(device)
    text_embeddings = model.get_text_features(**inputs)
    # convert the embeddings to numpy array
    embedding_as_np = text_embeddings.cpu().detach().numpy()

    return embedding_as_np

def get_item_data(result, index, measure_column) :

    img_name = str(result['image_name'][index])

    # TODO: add code to get the original comment
    comment = str(result['comment'][index])
    cos_sim = result[measure_column][index]

    return (img_name, comment, cos_sim)

def get_top_N_images(query,
                     data,
                     model, tokenizer,
                     device,
                     top_K=4) :

    query_vect = get_single_text_embedding(query, 
                                            model, tokenizer, 
                                            device)

    # Relevant columns
    relevant_cols = ["comment", "image_name", "cos_sim"]

    # Run similarity Search
    data["cos_sim"] = data["text_embeddings"].apply(lambda x: cosine_similarity(query_vect, x))# line 17
    data["cos_sim"] = data["cos_sim"].apply(lambda x: x[0][0])

    data_sorted = data.sort_values(by='cos_sim', ascending=False)
    non_repeated_images = ~data_sorted["image_name"].duplicated()
    most_similar_articles = data_sorted[non_repeated_images].head(top_K)

    """
    Retrieve top_K (4 is default value) articles similar to the query
    """

    result_df = most_similar_articles[relevant_cols].reset_index()

    return [get_item_data(result_df, i, 'cos_sim') for i in range(len(result_df))]


###### with faiss ###########

import faiss
import numpy as np

def faiss_add_index_cos(df, column):

    # Get the embeddings from the specified column
    embeddings = np.vstack(df[column].values).astype(np.float32)  # Convert to float32
          
    # Create an index
    index = faiss.IndexFlatIP(embeddings.shape[1])
    faiss.normalize_L2(embeddings) 

    index.train(embeddings)

    # Add the embeddings to the index
    index.add(embeddings)

    # Return the index
    return index


def faiss_get_top_N_images(query,
                           data,
                           model, tokenizer,
                           device,
                           top_K=4) :

    query_vect = get_single_text_embedding(query, 
                                          model, tokenizer, 
                                          device)
    # Relevant columns
    relevant_cols = ["comment", "image_name"]

    #faiss search with cos similarity
    index = faiss_add_index_cos(data, column="text_embeddings")

    faiss.normalize_L2(query_vect)
    D, I = index.search(query_vect, len(data))

    data_sorted = data.iloc[I.flatten()]

    non_repeated_images = ~data_sorted["image_name"].duplicated()
    most_similar_articles = data_sorted[non_repeated_images].head(top_K)

    result_df = most_similar_articles[relevant_cols].reset_index()
    D = D.reshape(-1,1)[:top_K]
    result_df = pd.concat([result_df, pd.DataFrame(D, columns=['similarity'])], axis=1)
    return [get_item_data(result_df, i, 'similarity') for i in range(len(result_df))]

	from transformers import CLIPModel, CLIPTokenizer
	from sklearn.metrics.pairwise import cosine_similarity
	import faiss

	from dataframe import *

	def get_model_info(model_ID, device):
	# Save the model to device
	model = CLIPModel.from_pretrained(model_ID).to(device)

	# Get the tokenizer
	tokenizer = CLIPTokenizer.from_pretrained(model_ID)

	# Return model, processor & tokenizer
	return model, tokenizer


	def get_single_text_embedding(text, model, tokenizer, device):
	inputs = tokenizer(text, return_tensors = "pt", max_length=77, truncation=True).to(device)
	text_embeddings = model.get_text_features(**inputs)
	# convert the embeddings to numpy array
	embedding_as_np = text_embeddings.cpu().detach().numpy()

	return embedding_as_np

	def get_item_data(result, index, measure_column) :

	img_name = str(result['image_name'][index])

	# TODO: add code to get the original comment
	comment = str(result['comment'][index])
	cos_sim = result[measure_column][index]

	return (img_name, comment, cos_sim)

	def get_top_N_images(query,
	data,
	model, tokenizer,
	device,
	top_K=4) :

	query_vect = get_single_text_embedding(query,
	model, tokenizer,
	device)

	# Relevant columns
	relevant_cols = ["comment", "image_name", "cos_sim"]

	# Run similarity Search
	data["cos_sim"] = data["text_embeddings"].apply(lambda x: cosine_similarity(query_vect, x))# line 17
	data["cos_sim"] = data["cos_sim"].apply(lambda x: x[0][0])

	data_sorted = data.sort_values(by='cos_sim', ascending=False)
	non_repeated_images = ~data_sorted["image_name"].duplicated()
	most_similar_articles = data_sorted[non_repeated_images].head(top_K)

	"""
	Retrieve top_K (4 is default value) articles similar to the query
	"""

	result_df = most_similar_articles[relevant_cols].reset_index()

	return [get_item_data(result_df, i, 'cos_sim') for i in range(len(result_df))]


	###### with faiss ###########

	import faiss
	import numpy as np

	def faiss_add_index_cos(df, column):

	# Get the embeddings from the specified column
	embeddings = np.vstack(df[column].values).astype(np.float32) # Convert to float32

	# Create an index
	index = faiss.IndexFlatIP(embeddings.shape[1])
	faiss.normalize_L2(embeddings)

	index.train(embeddings)

	# Add the embeddings to the index
	index.add(embeddings)

	# Return the index
	return index


	def faiss_get_top_N_images(query,
	data,
	model, tokenizer,
	device,
	top_K=4) :

	query_vect = get_single_text_embedding(query,
	model, tokenizer,
	device)
	# Relevant columns
	relevant_cols = ["comment", "image_name"]

	#faiss search with cos similarity
	index = faiss_add_index_cos(data, column="text_embeddings")

	faiss.normalize_L2(query_vect)
	D, I = index.search(query_vect, len(data))

	data_sorted = data.iloc[I.flatten()]

	non_repeated_images = ~data_sorted["image_name"].duplicated()
	most_similar_articles = data_sorted[non_repeated_images].head(top_K)

	result_df = most_similar_articles[relevant_cols].reset_index()
	D = D.reshape(-1,1)[:top_K]
	result_df = pd.concat([result_df, pd.DataFrame(D, columns=['similarity'])], axis=1)
	return [get_item_data(result_df, i, 'similarity') for i in range(len(result_df))]