| import json | |
| import time | |
| from openai import OpenAI | |
| with open('.openai-secret', 'r') as f: OPENAI_API_KEY = f.read() | |
| def read_dataset(): | |
| print("Reading dataset") | |
| json_file = open('updateddata.json') | |
| data = json.load(json_file) | |
| return data | |
| def main(): | |
| """ | |
| The purpse of this script is to take the articles json and add a field of the vector embeddings. | |
| There is a hardcoded sleep because the openAI API is rate limited to 1M tokens / min. | |
| In total, takes about an hour to run. | |
| """ | |
| data = read_dataset() | |
| print(len(data)) | |
| mini = 0 | |
| max = 750 | |
| client = OpenAI(api_key=OPENAI_API_KEY) | |
| newjson = [] | |
| while max <= len(data): | |
| print("------------") | |
| print("startind", mini, "endind", max) | |
| paper_subset = data[mini:max] | |
| abstract_list = [] | |
| for paper in paper_subset: | |
| abstract = paper['abstract'][0:2048] | |
| abstract = json.dumps(abstract, ensure_ascii=True) | |
| abstract_list.append(abstract) | |
| totallen = 0 | |
| for thinig in abstract_list: | |
| totallen+= len(thinig[0:2048]) | |
| print("numtokens:", totallen) | |
| abstract_list = [x.replace("\n"," ") for x in abstract_list] | |
| abstract_list = [x.strip() for x in abstract_list] | |
| res = client.embeddings.create( | |
| model="text-embedding-3-small", | |
| input=abstract_list, | |
| encoding_format="float" | |
| ) | |
| resdata = res.dict()['data'] | |
| print("Successful API call") | |
| for i in range(len(paper_subset)): | |
| paper_subset[i]['embed'] = resdata[i]['embedding'] | |
| print("Added embeds") | |
| newjson.append(paper_subset) | |
| if (max == len(data)): break | |
| time.sleep(61) | |
| mini += 750 | |
| max += 750 | |
| max = min(max, len(data)) | |
| with open("datawithembeds.json", 'w') as f: | |
| json.dump(newjson, f) | |
| if __name__ == '__main__': | |
| main() |