Update Cosine_distance.py

Cosine_distance.py

@@ -1,10 +1,9 @@
 import pandas as pd
-from gensim.models import Word2Vec
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 from nltk.tokenize import word_tokenize
-from keras.preprocessing.text import Tokenizer
-from keras.preprocessing.sequence import pad_sequences
+from transformers import AutoTokenizer, AutoModel
+import torch
 
 def recommend_papers(positive_csv_path, unlabelled_csv_path, n):
     # Load positive labelled data and unlabelled labelled data
@@ -15,37 +14,25 @@ def recommend_papers(positive_csv_path, unlabelled_csv_path, n):
     positive_labelled_info['text'] = positive_labelled_info['title'] + ' ' + positive_labelled_info['abstract']
     unlabelled_labelled['text'] = unlabelled_labelled['title'] + ' ' + unlabelled_labelled['abstract']
 
-    # Preprocess text and tokenize
-    positive_labelled_info['tokenized_text'] = positive_labelled_info['text'].apply(word_tokenize)
-    unlabelled_labelled['tokenized_text'] = unlabelled_labelled['text'].apply(word_tokenize)
+    # Tokenize text using the pre-trained Sentence Transformer model
+    tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+    positive_labelled_info['input_ids'] = positive_labelled_info['text'].apply(lambda x: tokenizer.encode(x, add_special_tokens=True))
+    unlabelled_labelled['input_ids'] = unlabelled_labelled['text'].apply(lambda x: tokenizer.encode(x, add_special_tokens=True))
 
-    # Train Word2Vec model on the tokenized text of positive labelled data
-    word2vec_model = Word2Vec(sentences=positive_labelled_info['tokenized_text'], vector_size=100, window=5, min_count=1, workers=4)
+    # Load the pre-trained Sentence Transformer model
+    model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
 
-    # Initialize tokenizer
-    tokenizer = Tokenizer()
-    tokenizer.fit_on_texts(positive_labelled_info['text'])
+    # Compute embeddings for positive labelled data
+    positive_labelled_info['embedding'] = positive_labelled_info['input_ids'].apply(lambda x: model(torch.tensor([x]))[0])
 
-    # Convert text to sequences of integer indices
-    positive_sequences = tokenizer.texts_to_sequences(positive_labelled_info['text'])
-    unlabelled_sequences = tokenizer.texts_to_sequences(unlabelled_labelled['text'])
-
-    # Pad or truncate sequences to a fixed length
-    max_length = 50  # Adjust as needed
-    padded_positive_sequences = pad_sequences(positive_sequences, maxlen=max_length, padding='post', truncating='post')
-    padded_unlabelled_sequences = pad_sequences(unlabelled_sequences, maxlen=max_length, padding='post', truncating='post')
-
-    # Compute centroid of positive examples' word embeddings
-    centroid = np.mean([word2vec_model.wv[token] for token in tokenizer.index_word.values() if token in word2vec_model.wv], axis=0)
+    # Compute centroid of positive examples' embeddings
+    centroid = np.mean(positive_labelled_info['embedding'].tolist(), axis=0)
 
     # Compute similarity scores for each document in the unlabelled labelled data
     similarity_scores = []
-    for sequence in padded_unlabelled_sequences:
-        embeddings = [word2vec_model.wv[token] for token in sequence if token in word2vec_model.wv]
-        if len(embeddings) > 0:
-            similarity_score = cosine_similarity([centroid], [np.mean(embeddings, axis=0)])[0][0]
-        else:
-            similarity_score = 0  # Default similarity score for empty embeddings
+    for input_ids in unlabelled_labelled['input_ids']:
+        embedding = model(torch.tensor([input_ids]))[0].detach().numpy()
+        similarity_score = cosine_similarity([centroid], [embedding])[0][0]
         similarity_scores.append(similarity_score)
 
     # Add similarity scores to unlabelled labelled data