-
Notifications
You must be signed in to change notification settings - Fork 0
/
embeddings_dataframe.py
89 lines (77 loc) · 2.37 KB
/
embeddings_dataframe.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import pandas as pd
import numpy as np
def sort_pmids(pmids):
"""
Sorts the pmids in the TREC or RELISH dataset into ascending
order.
Parameters
----------
pmids: List[str]
A list of all pubmed ids in the corpus.
Returns
-------
pmids: List[str]
A list of all pubmed ids in the corpus sorted
in the ascending order.
"""
# Convert all values in the list "pmids" from type 'str' to 'int'
pmids = list(map(int, pmids))
# Sort pmids in ascending order
pmids.sort()
return pmids
def get_embeddings(embeddings_path, pmids):
"""
Get a list of embeddings against the sorted pubmed ids.
Parameters
----------
embeddings_path: str
Path to the directory where the embeddings are
saved.
pmids: List[str]
A list of all pubmed ids in the corpus sorted
in the ascending order.
Returns
-------
embeddings_list: List[ndarray]
A list of embeddings (against the sorted
pmids).
"""
embeddings_list = []
for pmid in pmids:
# Load the embeddings
file_name = embeddings_path + str(pmid) + ".npy"
embeddings = np.load(file_name)
# Append the embeddings to the "embeddings_list"
embeddings_list.append(embeddings)
return embeddings_list
def save_dataframe(pmids, embeddings_list, output_file):
"""
Creates a 2-column Pandas dataframe with the pmids
and their embeddings. Saves the dataframe in a
Pickle Python format.
Parameters
----------
pmids: List[str]
A list of all pubmed ids in the corpus sorted
in the ascending order.
embeddings_list: List[ndarray]
A list of embeddings (against the sorted
pmids).
output_file: File path to save the Pandas dataframe.
"""
# Create pandas dataframe with pmids and embeddings
dict = {'pmids': pmids, 'embeddings': embeddings_list}
df = pd.DataFrame(dict)
# Save dataframe to Python pickle format
df.to_pickle(output_file)
def load_dataframe(file_path):
"""
Loads a 2-column Pandas dataframe with the pmids
and their embeddings, from a Pickle file.
Parameters
----------
file_path: Path where the Pickle file is saved.
"""
# Load dataframe
df = pd.read_pickle(file_path)
print (df)