Update func.py

func.py

@@ -12,7 +12,8 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 
 def filter_by_ganre(df: pd.DataFrame, ganre_list: list):
     filtered_df = df[df['ganres'].apply(lambda x: any(g in ganre_list for g in(x)))]
-    return filtered_df
+    filt_ind = filtered_df.index.to_list()
+    return filt_ind
 
 def mean_pooling(model_output, attention_mask):
     token_embeddings = model_output['last_hidden_state']
@@ -21,7 +22,7 @@ def mean_pooling(model_output, attention_mask):
     sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
     return sum_embeddings / sum_mask
 
-def recommendation(df: pd.DataFrame, embeddings:np.array, user_text: str, n=10):
+def recommendation(filt_ind: list, embeddings: np.array, user_text: str, n=10):
     token_user_text = tokenizer(user_text, return_tensors='pt', padding='max_length', truncation=True, max_length=512)
     user_embeddings = torch.Tensor().to(device)
     model.to(device)
@@ -31,7 +32,7 @@ def recommendation(df: pd.DataFrame, embeddings:np.array, user_text: str, n=10):
         outputs = model(**batch)
         user_embeddings = torch.cat([user_embeddings, mean_pooling(outputs, batch['attention_mask'])])
         user_embeddings = user_embeddings.cpu().numpy()
-    cosine_similarities = cosine_similarity(embeddings, user_embeddings.reshape(1, -1))
+    cosine_similarities = cosine_similarity(embeddings[filt_ind], user_embeddings.reshape(1, -1))
     df_res = pd.DataFrame(cosine_similarities.ravel(), columns=['cos_sim']).sort_values('cos_sim', ascending=False)
     dict_topn = df_res.iloc[:n, :].cos_sim.to_dict()
     return dict_topn