import pandas as pd import numpy as np import torch from transformers import BertModel, BertTokenizer from sklearn.metrics.pairwise import cosine_similarity tokenizer = BertTokenizer.from_pretrained("DeepPavlov/rubert-base-cased-sentence") model = BertModel.from_pretrained("DeepPavlov/rubert-base-cased-sentence", output_hidden_states = True) device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") 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)))] filt_ind = filtered_df.index.to_list() return filt_ind def mean_pooling(model_output, attention_mask): token_embeddings = model_output['last_hidden_state'] input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9) return sum_embeddings / sum_mask 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) model.eval() with torch.no_grad(): batch = {k: v.to(device) for k, v in token_user_text.items()} 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[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