# Import the load_dotenv function
from dotenv import load_dotenv
import os
import json
from langchain.docstore.document import Document
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
import json
import os
from langchain.docstore.document import Document
embeddings = HuggingFaceEmbeddings(model_name="dmlls/all-mpnet-base-v2-negation")
class TrialDatabaseBuilder:
def __init__(self, json_directory, desired_fields, fields_to_concatenate):
self.json_directory = json_directory
self.desired_fields = desired_fields
self.fields_to_concatenate = fields_to_concatenate
self.docs = []
self.ids = []
def load_json_files(self):
for filename in os.listdir(self.json_directory):
if filename.endswith('.json'):
file_path = os.path.join(self.json_directory, filename)
with open(file_path, 'r') as file:
json_data = json.load(file)
extracted_data = {field: json_data.get(field) for field in self.desired_fields}
self.ids.append(extracted_data["nct_id"])
metadata = {
"id": extracted_data.get("nct_id", ""),
"gender": extracted_data.get("gender", ""),
"condition": extracted_data.get("condition", ""),
"phase": extracted_data.get("phase", ""),
"minimum_age": extracted_data.get("minimum_age", ""),
"maximum_age": extracted_data.get("maximum_age", ""),
}
metadata = {k: v for k, v in metadata.items() if v is not None}
concatenated_string = ', '.join(str(extracted_data[field]) for field in self.fields_to_concatenate)
doc = Document(page_content=concatenated_string, metadata=metadata)
self.docs.append(doc)
def build_vectorstore(self):
vectorstore = Chroma.from_documents(self.docs, embeddings, persist_directory="../../data/db/", collection_name="trials")
vectorstore.persist()
vectorstore = None
class CriteriaDatabaseBuilder:
def __init__(self, json_directory, desired_fields):
self.json_directory = json_directory
self.desired_fields = desired_fields
self.docs = []
def load_json_files(self):
for filename in os.listdir(self.json_directory):
if filename.endswith('.json'):
file_path = os.path.join(self.json_directory, filename)
with open(file_path, 'r') as file:
json_data = json.load(file)
extracted_data = {field: json_data.get(field) for field in self.desired_fields}
eligibility_criteria = json_data.get("eligibility")
if eligibility_criteria is not None:
for index, criterion in enumerate(eligibility_criteria):
metadata = {
"nct_id": extracted_data['nct_id'],
"idx": index + 1,
}
metadata["criteria_type"] = criterion["entities_data"][0]["field"]
for i, entity in enumerate(criterion["entities_data"]):
for key, value in entity.items():
if key != "field":
metadata[f"{key}_{i + 1}"] = value
doc = Document(page_content=criterion["text"], metadata=metadata)
self.docs.append(doc)
def build_vectorstore(self):
vectorstore = Chroma.from_documents(self.docs, embeddings, persist_directory="../../data/db/", collection_name="criteria")
vectorstore.persist()
vectorstore = None
def main():
load_dotenv('../.env')
openai_access_key = os.getenv('OPENAI_ACCESS_KEY')
huggingface_token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
json_directory = '../../data/trials_jsons/'
desired_fields_criteria = ["nct_id", "eligibility"]
criteriadb_builder = CriteriaDatabaseBuilder(json_directory, desired_fields_criteria)
criteriadb_builder.load_json_files()
criteriadb_builder.build_vectorstore()
desired_fields_trials=["nct_id", "brief_title", "brief_summary", "condition", "gender", "minimum_age", "maximum_age", "phase"]
trialdb_builder = TrialDatabaseBuilder(json_directory, desired_fields_trials)
trialdb_builder.load_json_files()
trialdb_builder.build_vectorstore()
if __name__ == "__main__":
main()
Datasets
Models
