tokenizer = AutoTokenizer.from_pretrained('Twitter/twhin-bert-base')
model = AutoModel.from_pretrained('Twitter/twhin-bert-base')
def process_post_data(posts):
post_texts = [post.title + " " + post.content for post in posts]
inputs = tokenizer(post_texts, return_tensors="pt", padding=True, truncation=True)
return inputs
def get_recommendations(posts, liked_posts, disliked_posts, following_ids):
all_post_inputs = process_post_data(posts) if posts else None
liked_post_inputs = process_post_data(liked_posts) if liked_posts else None
disliked_post_inputs = process_post_data(disliked_posts) if disliked_posts else None
with torch.no_grad():
all_outputs = model(**all_post_inputs)
all_embeddings = all_outputs.last_hidden_state[:, 0, :]
if liked_post_inputs:
liked_outputs = model(**liked_post_inputs)
liked_embeddings = liked_outputs.last_hidden_state[:, 0, :]
else:
liked_embeddings = torch.zeros(1, model.config.hidden_size)
if disliked_post_inputs:
disliked_outputs = model(**disliked_post_inputs)
disliked_embeddings = disliked_outputs.last_hidden_state[:, 0, :]
else:
disliked_embeddings = torch.zeros(1, model.config.hidden_size)
similarities_liked = torch.matmul(all_embeddings, liked_embeddings.T)
similarities_disliked = torch.matmul(all_embeddings, disliked_embeddings.T)
recommendation_scores = torch.zeros(len(posts))
for i, post in enumerate(posts):
if post.author_id in following_ids:
recommendation_scores[i] += 50
recommendation_scores += similarities_liked.mean(dim=1) - similarities_disliked.mean(dim=1)
sorted_indices = torch.argsort(recommendation_scores, descending=True)
recommended_posts = [posts[i] for i in sorted_indices]
return recommended_posts