Update app.py

app.py

@@ -10,10 +10,9 @@ import requests
 
 model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-df = {0: pd.read_csv('data.csv'), 1: pd.read_csv('data2.csv')}
-embeddings = {0: np.load('embeddings2.npy'), 1: np.load('embeddings.npy')}
-for k in [0, 1]:
-  embeddings[k] = np.divide(embeddings[k], np.sqrt(np.sum(embeddings[k]**2, axis=1, keepdims=True)))
+df =  pd.read_csv('data2.csv')
+embeddings_npy = np.load('embeddings.npy')
+embeddings = np.divide(embeddings_npy, np.sqrt(np.sum(embeddings_npy**2, axis=1, keepdims=True)))
   
 def compute_text_embeddings(list_of_strings):
     inputs = processor(text=list_of_strings, return_tensors="pt", padding=True)
@@ -27,13 +26,10 @@ def download_img(path):
     return local_path
     
 def predict(query):
-    corpus = 'Movies'
     n_results=3
-
     text_embeddings = compute_text_embeddings([query]).detach().numpy()
-    k = 0 if corpus == 'Unsplash' else 1
-    results = np.argsort((embeddings[k]@text_embeddings.T)[:, 0])[-1:-n_results-1:-1]
-    paths = [download_img(df[k].iloc[i]['path']) for i in results]
+    results = np.argsort((embeddings@text_embeddings.T)[:, 0])[-1:-n_results-1:-1]
+    paths = [download_img(df.iloc[i]['path']) for i in results]
     print(paths)
     return paths