import os import time import asyncio from llama_index.core.query_engine import CitationQueryEngine from llama_index.core import VectorStoreIndex from llama_index.core import Settings from llama_index.embeddings.huggingface import HuggingFaceEmbedding from llama_index.llms.gemini import Gemini from llama_index.core.postprocessor import SimilarityPostprocessor from llama_index.core.storage.docstore import SimpleDocumentStore from llama_index.core import StorageContext, load_index_from_storage from llama_index.core.data_structs import Node from llama_index.core.schema import NodeWithScore import re import pandas as pd import gradio as gr import logging #Enable logging to see what's happening behind the scenes logging.basicConfig(level=logging.INFO) token_w = os.environ['token_w'] HF_TOKEN=os.environ['token_r'] API_KEY=os.environ["GOOGLE_API_KEY"] generation_config = { "temperature": 0, # "top_p": 1, # "top_k": 1, "max_output_tokens":8192, } safety_settings = [ { "category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE" }, ] llm = Gemini( model="models/gemini-1.5-flash-002", generation_config=generation_config, safety_settings=safety_settings, ) # Setup embedder embed_model_name = "BAAI/bge-small-en-v1.5" embed_model = HuggingFaceEmbedding(model_name=embed_model_name) Settings.llm = llm Settings.embed_model = embed_model # rebuild storage context storage_context = StorageContext.from_defaults(persist_dir="VectorStore") # load index index_persisted = load_index_from_storage(storage_context, index_id="vector_index") async def remove_ref(text): """Removes content after 'Reference Papers' (case-insensitive).""" split_text = re.split(r'\bReference Papers\b', text, flags=re.IGNORECASE) return split_text[0].strip() if len(split_text) > 1 else text.strip() async def clean_trial_text(text): """Removes intro text from references if present.""" sections, cleaned_sections, in_references = text.split('\n'), [], False has_intro_text, found_numbers, reference_title_index = False, False, -1 for i, line in enumerate(sections): if re.match(r'Reference Papers\s*$', line, re.IGNORECASE): in_references, reference_title_index = True, len(cleaned_sections) cleaned_sections.append(line) continue if in_references and not found_numbers: if re.match(r'\d+\.', line.strip()): found_numbers = True else: if line.strip(): has_intro_text = True cleaned_sections.append(line) continue if not in_references: cleaned_sections.append(line) if in_references and not has_intro_text and reference_title_index != -1: cleaned_sections.pop(reference_title_index) return '\n'.join(cleaned_sections).strip() async def get_criteria(study_information, top_k): """Fetches eligibility criteria and metadata for a study.""" criteria_response = await query_engine_get_study.aquery(f""" Based on the provided instructions and clinical trial information, generate the new eligibility criteria specific for clinical trial information. ### Instruction: Find suitable papers that are relevant or similar to the provided clinical trial information (### Clinical Trial Information). Prioritize the following topics when finding related studies: 1. Study Objectives 2. Study Design and Phases 3. Conditions 4. Intervention/Treatment Criteria Generation: As a clinical researcher, generate new eligibility criteria for the given clinical trial information. Analyze the information from all {top_k} related studies to generate new precise eligibility criteria. Ensure that the criteria are specific for the given clinical trial information (### Clinical Trial Information). Please follow the pattern of the output (### Pattern of the output). -------------------------------------------------- ### Clinical Trial Information {study_information} -------------------------------------------------- ### Pattern of the Output Inclusion Criteria 1. 2. ... Exclusion Criteria 1. 2. ... """) metadata_list = [source.node.get_metadata_str() for source in criteria_response.source_nodes] return criteria_response.response, metadata_list async def process_reference(metadata_list): """Formats metadata list into a numbered string.""" return "\n".join([f"{i + 1}. {meta}" for i, meta in enumerate(metadata_list)]) async def get_response(criteria, reference): """Processes eligibility criteria and updates references to match new numbering.""" response = await llm.acomplete(f""" ### Task Description: You are tasked with processing clinical trial metadata and eligibility criteria. The goal is to clean, reorder, and maintain consistency between the metadata and references used in eligibility criteria. ### Instructions: 1. Review the eligibility criteria provided, which include references to metadata numbers (e.g., [1], [2], etc.). Identify all reference numbers that are actually used in the criteria. 2. Remove metadata of reference papers (### Metadata of Reference Papers) that does not have a corresponding reference in the eligibility criteria. This will ensure only relevant references are kept. 3. Reorder the remaining metadata so that they are numbered sequentially, starting from 1. 4. Update the reference numbers in the eligibility criteria accordingly to reflect the new order. 5. Maintain Criteria Consistency: Ensure that the eligibility criteria remain exactly the same in terms of content, but the reference numbers are updated to match the new numbering of metadata. -------------------------------------------------- ### Eligibility Criteria {criteria} -------------------------------------------------- ### Metadata of Reference Papers {reference} -------------------------------------------------- ### Pattern of the Output Inclusion Criteria 1. 2. ... Exclusion Criteria 1. 2. ... Reference Papers 1.NCT ID: Study Name: Condition: Intervention/Treatment: 2.NCT ID: Study Name: Condition: Intervention/Treatment: . . .""") response_text = response.text return response_text async def extract_criteria(text): """Extracts inclusion and exclusion criteria from text.""" patterns = { "inclusion": r'Inclusion Criteria:?(.*?)(?=Exclusion Criteria)', "exclusion": r'Exclusion Criteria:?(.*?)(?=Reference Papers|\n\n\n)' } inclusion = re.search(patterns["inclusion"], text, re.DOTALL | re.IGNORECASE) exclusion = re.search(patterns["exclusion"], text, re.DOTALL | re.IGNORECASE) return ( "Inclusion Criteria:\n" + (inclusion.group(1).strip() if inclusion else "Not found") + "\n\n" + "Exclusion Criteria:\n" + (exclusion.group(1).strip() if exclusion else "Not found") ) async def run_function_on_text(top_k, study_obj, study_type, phase, purpose, allocation, intervention_model, Masking, conditions, interventions, location_countries, removed_location_countries): """Runs the main function to process study information and generate formatted output.""" query_engine_get_study = CitationQueryEngine.from_args( index_persisted, similarity_top_k=top_k, citation_chunk_size=2048, verbose=True, node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.8)], use_async=True ) study_information = f""" # Study Objectives/Description {study_obj} # Intervention {interventions} # Location - Location_Countries: {location_countries} - Removed Location: {removed_location_countries} # Conditions Cancer {conditions} # Study Design - Study Type: {study_type} - Phase: {phase} - Primary Purpose: {purpose} - Allocation: {allocation} - Interventional Model: {intervention_model} - Masking: None {Masking} """ criteria, metadata_list = await get_criteria(study_information, top_k) if criteria != "Empty Response": processed_ref = await process_reference(metadata_list) response = await get_response(criteria, processed_ref) combine_criteria = await extract_criteria(response) # Extract and format references pattern = r'Reference Papers\s*(.+)$' match = re.search(pattern, response, re.DOTALL | re.IGNORECASE) ext_ref = match.group(1) if match else "" split_ref = re.split(r'\n*\d+\.\s+', ext_ref)[1:] formatted_ref = [] for i, ref in enumerate(split_ref, 1): nct_id = re.search(r'NCT[_ ]ID: (NCT\d+)', ref) if not nct_id: nct_id = re.search(r'(NCT\d+)', ref) if not nct_id: continue study_name = re.search(r'Study[_ ]Name:?\s*(.*?)(?=\n|;|Condition:|Intervention/Treatment:|$)', ref, re.DOTALL) condition = re.search(r'Condition:?\s*(.*?)(?=\n|;|Intervention/Treatment:|$)', ref, re.DOTALL) intervention = re.search(r'Intervention/Treatment:?\s*(.*?)(?=\n|$)', ref, re.DOTALL) formatted_ref.append([ i, f'{nct_id.group(1)}', study_name.group(1).strip() if study_name else "", condition.group(1).strip() if condition else "", intervention.group(1).strip() if intervention else "" ]) else: combine_criteria, formatted_ref = "Empty Response", [] return combine_criteria, formatted_ref # # LLM.complete # complete_response = await llm.acomplete(f""" # Based on the provided instructions and clinical trial information, generate the new eligibility criteria by analyzing clinical trial information(### Clinical Trial Information). # ### Instruction: # Criteria generation: # As a clinical researcher, generate new eligibility criteria for given clinical trial information. # Ensure the criteria are clear, specific, and reasonable for a clinical research information. # Prioritize the following topics in clinical trial information.: # 1. Study Objectives # 2. Study Design and Phases # 3. Conditions # 4. Intervention/Treatment # Please follow the pattern of the output(### Pattern of the output). # -------------------------------------------------- # ### Clinical Trial Information # {study_information} # -------------------------------------------------- # ### Pattern of the output # Inclusion Criteria # 1. # 2. # . # . # . # Exclusion Criteria # 1. # 2. # . # . # . # """ # ) # combine_response = await llm.acomplete(f""" # Based on the provided instructions clinical, clinical trial information, and criteria information, generate the appropriate eligibility criteria for ### Clinical Trial Information by analyze clinical trial information(### Clinical Trial Information), criteria 1 (### Criteria 1) and criteria 2 (### Criteria 2). # ### Instruction: # Criteria generation: # As a clinical researcher, generate appropriate eligibility criteria by analyzing given information. # Ensure the criteria are clear, specific, and reasonable for a clinical research information(### Clinical Trial Information). # Prioritize the following topics in clinical trial information.: # 1. Study Objectives # 2. Study Design and Phases # 3. Conditions # 4. Intervention/Treatment # Do not generate redundant inclusion and exclusion criteria. For example, if a criterion is included in one set of inclusion or exclusion criteria, do not include it again. # Reference Papers generation: # Please give us NCT IDs and study names from the references list in ### Criteria 1. # Please follow the pattern of the output(### Pattern of the output). # -------------------------------------------------- # ### Clinical Trial Information # {study_information} # -------------------------------------------------- # ### Criteria 1 # {query_response} # -------------------------------------------------- # ### Criteria 2 # {complete_response} # -------------------------------------------------- # ### Pattern of the output # Inclusion Criteria # 1. # 2. # . # . # . # Exclusion Criteria # 1. # 2. # . # . # . # Reference Papers # 1.NCT ID: # Study Name: # Condition: # Intervention/Treatment: # 2.NCT ID: # Study Name: # Condition: # Intervention/Treatment: # . # . # . # """ # ) # return query_response # return query_response,complete_response,combine_response # Place holder place_holder = f"""Study Objectives The purpose of this study is to evaluate the safety, tolerance and efficacy of Liposomal Paclitaxel With Nedaplatin as First-line in patients with Advanced or Recurrent Esophageal Carcinoma Conditions: Esophageal Carcinoma Intervention / Treatment: DRUG: Liposomal Paclitaxel, DRUG: Nedaplatin Location: China Study Design and Phases Study Type: INTERVENTIONAL Phase: PHASE2 Primary Purpose: TREATMENT Allocation: NA Interventional Model: SINGLE_GROUP Masking: NONE """ objective_place_holder = f"""Example: The purpose of this study is to evaluate the safety, tolerance and efficacy of Liposomal Paclitaxel With Nedaplatin as First-line in patients with Advanced or Recurrent Esophageal Carcinoma """ conditions_place_holder = f"""Example: Esophageal Carcinoma """ interventions_place_holder = f"""Example: - Drug: irinotecan hydrochloride - Given IV - Other Names: - Campto - Camptosar - CPT-11 - irinotecan - U-101440E - Drug: Amoxicillin hydrate - Amoxicillin hydrate (potency) - Biological: Pneumococcal Vaccine - Subcutaneously on Day 0 - Other Names: - Prevnar - Drug: Doxorubicin, Cotrimoxazole, Carboplatin, Ifosfamide - Drug: Irinotecan - Irinotecan will be administered at a dose of 180mg/m2 IV over 90 minutes on day 21 every 42 days. - Other Names: - CAMPTOSAR™ - Drug: Placeblo - Placebo tablet """ with gr.Blocks() as demo: # Study description with gr.Row(): gr.Markdown("# Research Information"), with gr.Row(): study_obj_box = gr.Textbox( label="Study Objective / Study Description", placeholder=objective_place_holder, lines=10) # Conditions with gr.Row(): gr.Markdown("# Conditions"), with gr.Row(): conditions_box = gr.Textbox( label="Conditions / Disease", info="Primary Disease or Condition of Cancer Being Studied in the Trial, or the Focus of the Study", placeholder=conditions_place_holder, ) #Interventions with gr.Row(): gr.Markdown("# Interventions / Drugs"), with gr.Row(): intervention_box = gr.Textbox( label="Intervention type", info="A process or action studied in a clinical trial, including drugs, devices, procedures, vaccines, or noninvasive approaches.", placeholder=interventions_place_holder, # lines=5, ) # Study Design with gr.Row(): gr.Markdown("# Study Design"), with gr.Column(): study_type_box = gr.Radio( ["Expanded Access", "Interventional", "Observational"], label="Study Type", ) phase_box= gr.CheckboxGroup( ["Not Applicable", "Early Phase 1", "Phase 1", "Phase 2", "Phase 3", "Phase 4"], label="Phase" ) purpose_box = gr.Radio( ["Treatment", "Prevention", "Diagnostic", "Educational/Counseling/Training", "Supportive Care", "Screening", "Health Services Research", "Basic Science", "Device Feasibility", "Other"], label="Primary Purpose" ) allocation_box = gr.Radio( ["Randomized", "Non-Randomized", "N/A"], label="Allocation" ) intervention_model_box = gr.Radio( ["Parallel", "Single-Group", "Crossover", "Factorial", "Sequential"], label="Interventional Model" ) masking_box = gr.Radio( ["None (Open Label)", "Single", "Double", "Triple", "Quadruple"], label="Masking" ) #Location with gr.Row(): gr.Markdown("# Location"), with gr.Column(): location_box = gr.Textbox( label="Location (Countries)", ) removed_location_box = gr.Textbox( label="Removed Location (Countries)", ) # Reference paper with gr.Row(): gr.Markdown("# Reference paper"), with gr.Row(): top_k_box = gr.Slider( label="Amount of reference paper", value=10, minimum=0, maximum=30, step=1, ) # Submit & Clear with gr.Row(): submit_button = gr.Button("Submit") clear_button = gr.Button("Clear") # Output with gr.Row(): gr.Markdown("# Eligibility Criteria Generation"), with gr.Row(): with gr.Column(): base_box = gr.Textbox( label="Response", lines=15, interactive=False) with gr.Row(): ref_table = gr.Dataframe( label="Reference", headers=["No.",'Link', 'Study name', 'Intervention', 'Condition'], datatype=["markdown","html","markdown", "markdown","markdown"], wrap=True, interactive=False) # with gr.Column(): # rag_box = gr.Textbox( # label="Response 2", # lines=15, # interactive=False) # with gr.Column(): # combine_box = gr.Textbox( # label="Response 3", # lines=15, # interactive=False) with gr.Row(): regenerate_button = gr.Button("Regenerate") inputs_information = [top_k_box, study_obj_box, study_type_box, phase_box, purpose_box, allocation_box, intervention_model_box, masking_box, conditions_box, intervention_box, location_box, removed_location_box] outputs_information = [base_box,ref_table] # outputs_information = [base_box, rag_box,combine_box] submit_button.click( run_function_on_text, inputs=inputs_information, outputs=outputs_information ) regenerate_button.click( run_function_on_text, inputs=inputs_information, outputs=outputs_information ) clear_button.click(lambda : [None] * len(inputs_information), outputs=inputs_information) # with gr.Row(): # selected_response = gr.Radio( # choices=[ # "Response 1", # "Response 2", # "Response 3", # "All responses are equally good", # "Neither response is satisfactory" # ], # label="Select the best response" # ) # with gr.Row(): # flag_button = gr.Button("Flag Selected Response") # #Flagging # dataset_name = "ravistech/feedback-demo-space" # hf_writer = gr.HuggingFaceDatasetSaver(hf_token=token_w, dataset_name=dataset_name, private=True) # hf_writer.setup([selected_response, study_obj_box, study_type_box, phase_box, purpose_box, allocation_box, intervention_model_box, masking_box, conditions_box, intervention_box, location_box, removed_location_box, top_k_box, base_box, rag_box, combine_box],dataset_name) # flag_button.click(lambda *args: hf_writer.flag(list(args)), [selected_response, study_obj_box, study_type_box, phase_box, purpose_box, allocation_box, intervention_model_box, masking_box, conditions_box, intervention_box, location_box, removed_location_box, top_k_box, base_box, rag_box, combine_box], None, preprocess=False) #Clear all with gr.Row(): clear_all_button = gr.Button("Clear All") # flag_response = [selected_response] all_information = inputs_information + outputs_information #+ flag_response clear_all_button.click(lambda : [None] * len(all_information), outputs=all_information) if __name__ == "__main__": demo.launch(debug=True)