Update app.py

app.py

@@ -1,203 +1,227 @@
 import os
-token_r=os.environ['token_r']
-token_r=os.environ['token_r_nt']
-token_w=os.environ['token_w']
-token_w_feedback=os.environ['token_w_feedback']
-es_cloud_id=os.environ['es_cloud_id']
-feedback_store=os.environ['feedback_store']
-es_password=os.environ['es_password']
-
-from llama_index.core import Settings
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.vector_stores.elasticsearch import ElasticsearchStore
+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
 
-import torch
-from transformers import AutoTokenizer
-from llama_index.llms.huggingface import HuggingFaceLLM
-from transformers import BitsAndBytesConfig
-import asyncio
-
-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)
-
-model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
-
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    token=token_r,
-)
-
-stopping_ids = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>"),
+HF_TOKEN=os.environ["HUGGINGFACE_TOKEN"]
+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"
+  },
 ]
 
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
+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)
 
-# Get the model
-llm = HuggingFaceLLM(
-    model_name=model_name,
-    tokenizer_name=model_name,
-    model_kwargs={
-        "token": token_r,
-        "quantization_config": quantization_config
-    },
-    context_window=8191,
-    max_new_tokens=2048,
-    generate_kwargs={
-        # "do_sample": True,
-        # "temperature": 0.1,
-        # "top_p": 0.9,
-        'repetition_penalty': 1.175,
-        # 'early_stopping': True
-    },
-    stopping_ids=stopping_ids,
-)
-
-# bge embedding model
-embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
-Settings.embed_model = embed_model
-
-# Llama-3-8B-Instruct model
 Settings.llm = llm
+Settings.embed_model = embed_model
 
-# Get data from Elasticsearch
-
-es_vector_store = ElasticsearchStore(
-    index_name="train_criteria_index",
-    es_cloud_id=es_cloud_id,
-    es_user="elastic",
-    es_password=es_password,
-)
+#to-do: get data
 
-index_es = VectorStoreIndex.from_vector_store(es_vector_store)
+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):
 
-query_engine_get_study = CitationQueryEngine.from_args(
-    index_es,
-    similarity_top_k=5,
+  # Set up query engine
+  query_engine_get_study = CitationQueryEngine.from_args(
+    index_persisted,
+    similarity_top_k=top_k,
     citation_chunk_size=2048,
-    node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.8)],
-    use_async=True,    
     verbose=True,
-)
+    node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.8)],
+    use_async=True
+  )
+
+  #Build prompt
+  study_information = f"""
+  #Study Objectives/Study 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}
+  """
+
+  # Query
+  query_response = await query_engine_get_study.aquery(f"""
+      Based on the provided instructions and clinical trial information, generate the new eligibility criteria by analyzing the related studies and clinical trial information.
+      ### Instruction:
+      Find suitable papers that have relevant or similar to the 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 given clinical trial information.
+      Analyze the information from related studies for more precise new eligibility criteria generation.
+      Ensure the criteria are clear, specific, and reasonable for a clinical research information.
+      Do not generate [<number of citation>].
+
+      Reference Papers generation:
+      Please give us NCT IDs and study names for {top_k} used papers.
+
+      Please follows 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.
+      .
+      .
+      .
+
+      Reference Papers
+      1.NCT ID:
+        Study Name:
+        Condition:
+        Intervention/Treatment:
+      2.NCT ID:
+        Study Name:
+        Condition:
+        Intervention/Treatment:
+      .
+      .
+      .
+
+
+      """
+    )
 
-# # Regex patterns for output cleaning 
-# inclusion_pattern = r'Inclusion Criteria(.*?)(?=\*\*Exclusion Criteria)'
-# exclusion_pattern = r'Exclusion Criteria(.*?)(?=\*\*Reference Papers)'
-# reference_pattern = r'Reference Papers(.*)'
-
-# # Flags for multiline and dotall
-# flags = re.MULTILINE | re.DOTALL
-
-# # Function to clean text
-# def clean_text(text):
-#     # Remove asterisks
-#     cleaned = re.sub(r'\*', '', text)
-#     # # Remove extra whitespace
-#     # cleaned = re.sub(r'\s+', ' ', cleaned)
-#     # Remove leading/trailing whitespace
-#     return cleaned.strip()
-
-async def get_prompt(text):
-    studies_response = await query_engine_get_study.aquery(f"""
-    Based on the provided instructions and clinical trial information, delimited by ### Clinical Trial Information, which study has similar study information?
-    Ensure the studies are relevant and have similar study information. Prioritize the following topics when finding related studies:
-    1. Conditions
-    2. Intervention/Treatment
-    3. Study Objectives
-    4. Study Design and Phases    
-    
-    ### Clinical Trial Information:
-    {text}
-    """)
-    
-    study_ref=[]
-    metadata_list = []
-    for source in studies_response.source_nodes:
-        ref = source.node.get_text()
-        study_ref.append(ref)
-        meta_data = source.node.get_metadata_str()
-        metadata_list.append(meta_data)
-    
-    # return stream_
-    criteria_response = await llm.acomplete(f"""
-    Based on the provided instructions and clinical trial information, generate the eligibility criteria for the study.
-    
-    ## Instruction:
-    As a clinical researcher, generate new eligibility criteria for clinical research based on the given clinical trial information.
-    Please analyze the given clinical trial information, delimited by ### Clinical Trial Information, and analyze the information from the following papers, delimited by ### Related data, for more precise new eligibility criteria generation. Do not copy the ### Related data.
-    Please provide me NCT IDs and study names using information from the following papers, delimited by ### Reference Papers.
-    Please follow the pattern of the output, delimited by ### Pattern of the output.
-    Ensure the criteria are clear, specific, and suitable for a clinical research information.
-    Do not generate any other topic, explain text, python code, or any step of thought except the pattern of the output.
-
-    --------------------------------------------------
-    ### Clinical Trial Information
-    {text}
-    --------------------------------------------------
-    ### Related data
-    {study_ref}
-    --------------------------------------------------
-    ### Reference Papers
-    {metadata_list}
-    --------------------------------------------------
-    ### Pattern of the output
-    Inclusion Criteria
-    1.
-    2.
-
-    Exclusion Criteria
-    1.
-    2.
-
-    Reference Papers
-    1. NCT ID:
-      Study Name:
-    2. NCT ID:
-      Study Name:
-    3. NCT ID:
-      Study Name:
-      
-    """
+  # 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 follows 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 and criteria information, generate the new eligibility criteria by analyzing Criteria 1 (### Criteria 1) and Criteria 2 (### Criteria 2).
+      ### Instruction:
+      Criteria generation:
+      As a clinical researcher, generate eligibility criteria for given 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 follows 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.
+      .
+      .
+      .
+
+
+      """
     )
-    
-    # # Extract and clean data
-    # inclusion_criteria = clean_text(re.search(inclusion_pattern, text, flags).group(1))
-    # exclusion_criteria = clean_text(re.search(exclusion_pattern, text, flags).group(1))
-    # reference_papers = clean_text(re.search(reference_pattern, text, flags).group(1))
-    
-    # # Combine all sections into one formatted string
-    # criteria_response = f"""Inclusion Criteria:
-    # {inclusion_criteria}
-    
-    # Exclusion Criteria:
-    # {exclusion_criteria}
-    
-    # Reference Papers:
-    # {reference_papers}"""
-
-    # return criteria_response
-
-    return criteria_response.text
-    # for chunk in criteria_response:
-    #     yield chunk
 
+  return query_response,complete_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
 
@@ -216,46 +240,217 @@ TREATMENT Allocation: NA
 Interventional Model: SINGLE_GROUP Masking: NONE
 """
 
-prefilled_value = """Study Objectives
-[Brief Summary and/or Detailed Description]
+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
+"""
 
-Conditions: [Disease]
+interventions_place_holder = f"""Example:
+- Drug: irinotecan hydrochloride
+   - Given IV
+   - Other Names:
+       - Campto
+       - Camptosar
+       - CPT-11
+       - irinotecan
+       - U-101440E
 
-Intervention / Treatment
-[DRUGs]
+- Drug: Amoxicillin hydrate
+   - Amoxicillin hydrate (potency)
 
-Location
-[Location]
+- Procedure: Stem cell transplant
+   - See Detailed Description section for details of treatment interventions.
 
-Study Design and Phases
-Study Type:
-Phase:
-Primary Purpose:
-Allocation:
-Interventional Model:
-Masking:"""
-
-hf_writer = gr.HuggingFaceDatasetSaver(hf_token=token_w_feedback, dataset_name="nttwt1597/criteria-feedback-demo", private=True)
-
-prompt_box = gr.Textbox(
-    lines=10,
-    label=" Research Information",
-    # placeholder=place_holder,
-    value=prefilled_value,
-)
+- Biological: Pneumococcal Vaccine
+   - Subcutaneously on Day 0
+   - Other Names:
+       - Prevnar
 
-output_box = gr.Textbox(
-    lines=10,
-    label="Eligible criteria Criteria",
-)
+- 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
+"""
+
+prefilled_value = f"""Study Objectives The purpose of this study is to find out if the combination of docetaxel and capecitabine can shrink the size of breast tumors and preserve the breast. Conditions: Breast Cancer Intervention / Treatment: DRUG: Docetaxel, DRUG: Capecitabine Location: United States Study Design and Phases Study Type: INTERVENTIONAL Phase: PHASE2 Primary Purpose: TREATMENT Allocation: RANDOMIZED Interventional Model: PARALLEL Masking: NONE"""
+
+custom_css = """
+.gradio-container {
+    font-family: 'Roboto', sans-serif;
+}
+.main-header {
+    text-align: center;
+    color: #4a4a4a;
+    margin-bottom: 2rem;
+}
+.tab-header {
+    font-size: 1.2rem;
+    font-weight: bold;
+    margin-bottom: 1rem;
+}
+.custom-chatbot {
+    border-radius: 10px;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+}
+.custom-button {
+    background-color: #3498db;
+    color: white;
+    border: none;
+    padding: 10px 20px;
+    border-radius: 5px;
+    cursor: pointer;
+    transition: background-color 0.3s ease;
+}
+.custom-button:hover {
+    background-color: #2980b9;
+}
+"""
 
-demo = gr.Interface(
-    fn=get_prompt,
-    inputs=prompt_box,
-    outputs=output_box,
-    allow_flagging="manual",
-    flagging_options=["appropriate", "inappropriate", "incorrect",],
-    flagging_callback=hf_writer,
+# Define Gradio theme
+theme = gr.themes.Default(
+    primary_hue="zinc",
+    secondary_hue="red",
+    neutral_hue="neutral",
+    font=[gr.themes.GoogleFont('Roboto'), "sans-serif"]
 )
 
-demo.queue(max_size=20,default_concurrency_limit=5 ).launch(debug=True, share=True)
+with gr.Blocks() as demo:
+
+  with gr.Row():
+      gr.Markdown("# Reference paper"),
+  with gr.Row():
+      top_k_box = gr.Slider(
+          label="Amount of reference paper",
+          value=5,
+          minimum=0,
+          maximum=30,
+          step=1,
+          )
+  # Study description
+  with gr.Row():
+      gr.Markdown("# Research Information"),
+  with gr.Row():
+      study_obj_box = gr.Textbox(
+          label="Study Objective / Study Description", # Study description
+          placeholder=objective_place_holder,
+          lines=10)
+  # 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.Radio(
+          ["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"
+          )
+  # 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",
+          placeholder=interventions_place_holder,
+          )
+  #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)",
+          )
+
+  # 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 1",
+              lines=5,
+              interactive=False)
+      with gr.Column():
+          rag_box = gr.Textbox(
+              label="Response 2",
+              lines=5,
+              interactive=False)
+
+  submit_button.click(
+      run_function_on_text,
+      inputs=[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=[base_box, rag_box]
+  )
+
+  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]
+
+  clear_button.click(lambda : [None] * len(inputs_information), outputs=inputs_information)
+
+
+  with gr.Row():
+      selected_response = gr.Radio(
+          choices=[
+              "Response 1 is better",
+              "Response 2 is better",
+              "Both responses are equally good",
+              "Neither response is satisfactory"
+              ],
+          label="Select the better response"
+      )
+  with gr.Row():
+      flag_button = gr.Button("Flag Selected Response")
+
+  #Flagging
+  callback = gr.CSVLogger()
+  callback.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], "flagged_data_points")
+
+  flag_button.click(lambda *args: callback.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], None, preprocess=False)
+
+  with gr.Row():
+    clear_all_button = gr.Button("Clear All")
+
+  output_information = inputs_information + [base_box, rag_box, selected_response]
+  clear_all_button.click(lambda : [None] * len(output_information), outputs=output_information)
+
+if __name__ == "__main__":
+  demo.launch(debug=True)