File size: 1,270 Bytes
8df0b72 63d858e 8df0b72 a5cb4ef d760786 63d858e a5cb4ef 63d858e af858b4 63d858e af858b4 d760786 af858b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
import random
import torch
from dataframe import *
from model import *
def search(search_prompt : str) :
# Set the device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Define the model ID
model_ID = "openai/clip-vit-base-patch32"
# Get model, processor & tokenizer
model, tokenizer = get_model_info(model_ID, device)
image_data_df = get_image_data('data/output2.csv')
return get_top_N_images(search_prompt,
data = image_data_df,
model=model, tokenizer=tokenizer,
device = device,
top_K=4)
def searchWithFaiss(search_prompt : str) :
# Set the device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Define the model ID
model_ID = "openai/clip-vit-base-patch32"
# Get model, processor & tokenizer
model, tokenizer = get_model_info(model_ID, device)
image_data_df = get_image_data('data/output2.csv')
return faiss_get_top_N_images(search_prompt,
data = image_data_df,
model=model, tokenizer=tokenizer,
device = device,
top_K=4) |