from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer from sklearn.metrics.pairwise import cosine_similarity 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 processor processor = CLIPProcessor.from_pretrained(model_ID) # Get the tokenizer tokenizer = CLIPTokenizer.from_pretrained(model_ID) # Return model, processor & tokenizer return model, processor, 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 df_to_array(result_df) : return [str(result_df['image_name'][i]) for i in range(len(result_df))] def get_top_N_images(query, data, model, tokenizer, device, top_K=4, search_criterion="text"): # Text to image Search if (search_criterion.lower() == "text"): query_vect = get_single_text_embedding(query, model, tokenizer, device) # # Image to image Search # else: # query_vect = get_single_image_embedding(query) # Relevant columns revevant_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[revevant_cols].reset_index() return df_to_array(result_df)