Upload 8 files

app.py

@@ -1,50 +1,62 @@
 import gradio as gr
 import mtdna_backend
 import json
+from iterate3 import data_preprocess, model, pipeline
+import os
+import hashlib
+import threading
 # Gradio UI
+#stop_flag = gr.State(value=False)
+class StopFlag:
+    def __init__(self):
+        self.value = False
+global_stop_flag = StopFlag()  # Shared between run + stop
+
 with gr.Blocks() as interface:
     gr.Markdown("# 🧬 mtDNA Location Classifier (MVP)")
 
-    inputMode = gr.Radio(choices=["Single Accession", "Batch Input"], value="Single Accession", label="Choose Input Mode")
-
-    with gr.Group() as single_input_group:
-        single_accession = gr.Textbox(label="Enter Single Accession (e.g., KU131308)")
+    #inputMode = gr.Radio(choices=["Single Accession", "Batch Input"], value="Single Accession", label="Choose Input Mode")
+    user_email = gr.Textbox(label="📧 Your email (used to track free quota)")
+    usage_display = gr.Markdown("", visible=False)
 
-    with gr.Group(visible=False) as batch_input_group:
-        raw_text = gr.Textbox(label="🧬 Paste Accession Numbers (e.g., MF362736.1,MF362738.1,KU131308,MW291678)")
-        gr.HTML("""<a href="https://drive.google.com/file/d/1t-TFeIsGVu5Jh3CUZS-VE9jQWzNFCs_c/view?usp=sharing" download target="_blank">Download Example CSV Format</a>""")
-        gr.HTML("""<a href="https://docs.google.com/spreadsheets/d/1lKqPp17EfHsshJGZRWEpcNOZlGo3F5qU/edit?usp=sharing&ouid=112390323314156876153&rtpof=true&sd=true" download target="_blank">Download Example Excel Format</a>""")
-        file_upload = gr.File(label="📁 Or Upload CSV/Excel File", file_types=[".csv", ".xlsx"], interactive=True, elem_id="file-upload-box")
-        
+    # with gr.Group() as single_input_group:
+    #     single_accession = gr.Textbox(label="Enter Single Accession (e.g., KU131308)")
 
+    # with gr.Group(visible=False) as batch_input_group:
+    #     raw_text = gr.Textbox(label="🧬 Paste Accession Numbers (e.g., MF362736.1,MF362738.1,KU131308,MW291678)")
+    #     resume_file = gr.File(label="🗃️ Previously saved Excel output (optional)", file_types=[".xlsx"], interactive=True)
+    #     gr.HTML("""<a href="https://drive.google.com/file/d/1t-TFeIsGVu5Jh3CUZS-VE9jQWzNFCs_c/view?usp=sharing" download target="_blank">Download Example CSV Format</a>""")
+    #     gr.HTML("""<a href="https://docs.google.com/spreadsheets/d/1lKqPp17EfHsshJGZRWEpcNOZlGo3F5qU/edit?usp=sharing&ouid=112390323314156876153&rtpof=true&sd=true" download target="_blank">Download Example Excel Format</a>""")
+    #     file_upload = gr.File(label="📁 Or Upload CSV/Excel File", file_types=[".csv", ".xlsx"], interactive=True, elem_id="file-upload-box")
+    raw_text = gr.Textbox(label="🧚 Input Accession Number(s) (single (KU131308) or comma-separated (e.g., MF362736.1,MF362738.1,KU131308,MW291678))")
+    #resume_file = gr.File(label="🗃️ Previously saved Excel output (optional)", file_types=[".xlsx"], interactive=True)
+    gr.HTML("""<a href="https://docs.google.com/spreadsheets/d/1lKqPp17EfHsshJGZRWEpcNOZlGo3F5qU/edit?usp=sharing" download target="_blank">Download Example Excel Format</a>""")
+    file_upload = gr.File(label="📁 Or Upload CSV/Excel File", file_types=[".csv", ".xlsx"], interactive=True)
 
     with gr.Row():
         run_button = gr.Button("🔍 Submit and Classify")
+        stop_button = gr.Button("❌ Stop Batch", visible=True)
         reset_button = gr.Button("🔄 Reset")
 
     status = gr.Markdown(visible=False)
-
+    
     with gr.Group(visible=False) as results_group:
-      with gr.Accordion("Open to See the Result", open=False) as results:  
-          with gr.Row():
-              output_summary = gr.Markdown(elem_id="output-summary")
-              output_flag = gr.Markdown(elem_id="output-flag")
+      # with gr.Accordion("Open to See the Result", open=False) as results:  
+      #     with gr.Row():
+      #         output_summary = gr.Markdown(elem_id="output-summary")
+      #         output_flag = gr.Markdown(elem_id="output-flag")
           
-          gr.Markdown("---")
+      #     gr.Markdown("---")
 
       with gr.Accordion("Open to See the Output Table", open=False) as table_accordion:    
-          """output_table = gr.Dataframe(
-              headers=["Sample ID", "Technique", "Source", "Predicted Location", "Haplogroup", "Inferred Region", "Context Snippet"],
-              interactive=False,
-              row_count=(5, "dynamic")
-          )"""
           output_table = gr.HTML(render=True)
 
-
       with gr.Row():
           output_type = gr.Dropdown(choices=["Excel", "JSON", "TXT"], label="Select Output Format", value="Excel")
           download_button = gr.Button("⬇️ Download Output")
-      download_file = gr.File(label="Download File Here",visible=False)
+      #download_file = gr.File(label="Download File Here",visible=False)
+      download_file = gr.File(label="Download File Here", visible=False, interactive=True)
+      progress_box = gr.Textbox(label="Live Processing Log", lines=20, interactive=False)
 
       gr.Markdown("---")
 
@@ -56,47 +68,221 @@ with gr.Blocks() as interface:
       feedback_status = gr.Markdown()
 
     # Functions
-
-    def toggle_input_mode(mode):
-        if mode == "Single Accession":
-            return gr.update(visible=True), gr.update(visible=False)
-        else:
-            return gr.update(visible=False), gr.update(visible=True)
+    # def toggle_input_mode(mode):
+    #     if mode == "Single Accession":
+    #         return gr.update(visible=True), gr.update(visible=False)
+    #     else:
+    #         return gr.update(visible=False), gr.update(visible=True)
 
     def classify_with_loading():
         return gr.update(value="⏳ Please wait... processing...",visible=True)  # Show processing message
          
-    def classify_dynamic(single_accession, file, text, mode):
-        if mode == "Single Accession":
-            return classify_main(single_accession)  + (gr.update(visible=False),)
-        else:
-            #return summarize_batch(file, text) + (gr.update(visible=False),)  # Hide processing message
-            return classify_mulAcc(file, text) + (gr.update(visible=False),)  # Hide processing message
+    # def classify_dynamic(single_accession, file, text, resume, email, mode):
+    #     if mode == "Single Accession":
+    #         return classify_main(single_accession)  + (gr.update(visible=False),)
+    #     else:
+    #         #return summarize_batch(file, text) + (gr.update(visible=False),)  # Hide processing message
+    #         return classify_mulAcc(file, text, resume) + (gr.update(visible=False),)  # Hide processing message
+    # Logging helpers defined early to avoid NameError
+    
+
+    # def classify_dynamic(single_accession, file, text, resume, email, mode):
+    #   if mode == "Single Accession":
+    #       return classify_main(single_accession) + (gr.update(value="", visible=False),)
+    #   else:
+    #       return classify_mulAcc(file, text, resume, email, log_callback=real_time_logger, log_collector=log_collector)
 
     # for single accession
-    def classify_main(accession):
-        table, summary, labelAncient_Modern, explain_label = mtdna_backend.summarize_results(accession)
-        flag_output = f"### 🏺 Ancient/Modern Flag\n**{labelAncient_Modern}**\n\n_Explanation:_ {explain_label}"
-        return (
-            #table,
-            make_html_table(table),
-            summary,
-            flag_output,
-            gr.update(visible=True),
-            gr.update(visible=False)
-        )
-    # for batch accessions
-    def classify_mulAcc(file, text):
-        table, summary, flag_output, gr1, gr2 = mtdna_backend.summarize_batch(file, text)
-        #flag_output = f"### 🏺 Ancient/Modern Flag\n**{labelAncient_Modern}**\n\n_Explanation:_ {explain_label}"
-        return (
-            #table,
-            make_html_table(table),
-            summary,
-            flag_output,
-            gr.update(visible=True),
-            gr.update(visible=False)
-        )
+    # def classify_main(accession):
+    #     #table, summary, labelAncient_Modern, explain_label = mtdna_backend.summarize_results(accession)
+    #     table = mtdna_backend.summarize_results(accession)
+    #     #flag_output = f"### 🏺 Ancient/Modern Flag\n**{labelAncient_Modern}**\n\n_Explanation:_ {explain_label}"
+    #     return (
+    #         #table,
+    #         make_html_table(table),
+    #         # summary,
+    #         # flag_output,
+    #         gr.update(visible=True),
+    #         gr.update(visible=False),
+    #         gr.update(visible=False)
+    #     )
+    
+    #stop_flag = gr.State(value=False)
+    #stop_flag = StopFlag()
+
+    # def stop_batch(stop_flag):
+    #   stop_flag.value = True
+    #   return gr.update(value="❌ Stopping...", visible=True), stop_flag
+    def stop_batch():
+      global_stop_flag.value = True
+      return gr.update(value="❌ Stopping...", visible=True)
+
+    # def threaded_batch_runner(file, text, email):
+    #   global_stop_flag.value = False
+    #   log_lines = []
+
+    #   def update_log(line):
+    #       log_lines.append(line)
+    #       yield (
+    #           gr.update(visible=False),  # output_table (not yet)
+    #           gr.update(visible=False),  # results_group
+    #           gr.update(visible=False),  # download_file
+    #           gr.update(visible=False),  # usage_display
+    #           gr.update(value="⏳ Still processing...", visible=True),  # status
+    #           gr.update(value="\n".join(log_lines))  # progress_box
+    #       )
+
+    #   # Start a dummy update to say "Starting..."
+    #   yield from update_log("🚀 Starting batch processing...")
+
+    #   rows, file_path, count, final_log, warning = mtdna_backend.summarize_batch(
+    #       file=file,
+    #       raw_text=text,
+    #       resume_file=None,
+    #       user_email=email,
+    #       stop_flag=global_stop_flag,
+    #       yield_callback=lambda line: (yield from update_log(line))
+    #   )
+
+    #   html = make_html_table(rows)
+    #   file_update = gr.update(value=file_path, visible=True) if os.path.exists(file_path) else gr.update(visible=False)
+    #   usage_or_warning_text = f"**{count}** samples used by this email." if email.strip() else warning
+
+    #   yield (
+    #       html,
+    #       gr.update(visible=True),        # results_group
+    #       file_update,                    # download_file
+    #       gr.update(value=usage_or_warning_text, visible=True),
+    #       gr.update(value="✅ Done", visible=True),
+    #       gr.update(value=final_log)
+    #   )
+    
+    def threaded_batch_runner(file=None, text="", email=""):
+      print("📧 EMAIL RECEIVED:", email)
+      import tempfile
+      from mtdna_backend import (
+          extract_accessions_from_input,
+          summarize_results,
+          save_to_excel,
+          hash_user_id,
+          increment_usage,
+      )
+      import os
+
+      global_stop_flag.value = False  # reset stop flag
+
+      tmp_dir = tempfile.mkdtemp()
+      output_file_path = os.path.join(tmp_dir, "batch_output_live.xlsx")
+      limited_acc = 50 + (10 if email.strip() else 0)
+
+      # Step 1: Parse input
+      accessions, error = extract_accessions_from_input(file, text)
+      if error:
+          yield (
+              "",                          # output_table
+              gr.update(visible=False),   # results_group
+              gr.update(visible=False),   # download_file
+              "",                          # usage_display
+              "❌ Error",                  # status
+              str(error)                        # progress_box
+          )
+          return
+
+      total = len(accessions)
+      if total > limited_acc:
+          accessions = accessions[:limited_acc]
+          warning = f"⚠️ Only processing first {limited_acc} accessions."
+      else:
+          warning = f"✅ All {total} accessions will be processed."
+
+      all_rows = []
+      log_lines = []
+
+      # Step 2: Loop through accessions
+      for i, acc in enumerate(accessions):
+          if global_stop_flag.value:
+              log_lines.append(f"🛑 Stopped at {acc} ({i+1}/{total})")
+              usage_text = ""
+              if email.strip():
+                  # user_hash = hash_user_id(email)
+                  # usage_count = increment_usage(user_hash, len(all_rows))
+                  usage_count = increment_usage(email, len(all_rows))
+                  usage_text = f"**{usage_count}** samples used by this email. Ten more samples are added first (you now have 60 limited accessions), then wait we will contact you via this email."
+              else:
+                  usage_text = f"The limited accession is 50. The user has used {len(all_rows)}, and only {50-len(all_rows)} left."    
+              yield (
+                  make_html_table(all_rows),
+                  gr.update(visible=True),
+                  gr.update(value=output_file_path, visible=True),
+                  gr.update(value=usage_text, visible=True),
+                  "🛑 Stopped",
+                  "\n".join(log_lines)
+              )
+              return
+
+          log_lines.append(f"[{i+1}/{total}] Processing {acc}")
+          yield (
+              make_html_table(all_rows),
+              gr.update(visible=True),
+              gr.update(visible=False),
+              "",
+              "⏳ Processing...",
+              "\n".join(log_lines)
+          )
+
+          try:
+              rows = summarize_results(acc)
+              all_rows.extend(rows)
+              save_to_excel(all_rows, "", "", output_file_path, is_resume=False)
+              log_lines.append(f"✅ Processed {acc} ({i+1}/{total})")
+          except Exception as e:
+              log_lines.append(f"❌ Failed to process {acc}: {e}")
+
+          yield (
+              make_html_table(all_rows),
+              gr.update(visible=True),
+              gr.update(visible=False),
+              "",
+              "⏳ Processing...",
+              "\n".join(log_lines)
+          )
+
+      # Final update
+      usage_text = ""
+      if email.strip():
+          # user_hash = hash_user_id(email)
+          # usage_count = increment_usage(user_hash, len(all_rows))
+          usage_count = increment_usage(email, len(all_rows))
+          usage_text = f"**{usage_count}** samples used by this email. Ten more samples are added first (you now have 60 limited accessions), then wait we will contact you via this email."
+      else:
+          usage_text = f"The limited accession is 50. The user has used {len(all_rows)}, and only {50-len(all_rows)} left."
+      yield (
+          make_html_table(all_rows),
+          gr.update(visible=True),
+          gr.update(value=output_file_path, visible=True),
+          gr.update(value=usage_text, visible=True),
+          "✅ Done",
+          "\n".join(log_lines)
+      )
+
+    # def threaded_batch_runner(file=None, text="", email=""):
+    #   global_stop_flag.value = False
+
+    #   # Dummy test output that matches expected schema
+    #   return (
+    #       "<div>✅ Dummy output table</div>",   # HTML string
+    #       gr.update(visible=True),             # Group visibility
+    #       gr.update(visible=False),            # Download file
+    #       "**0** samples used.",               # Markdown
+    #       "✅ Done",                           # Status string
+    #       "Processing finished."               # Progress string
+    #   )
+
+    
+    # def classify_mulAcc(file, text, resume, email, log_callback=None, log_collector=None):
+    #     stop_flag.value = False
+    #     return threaded_batch_runner(file, text, resume, email, status, stop_flag, log_callback=log_callback, log_collector=log_collector)
+    
 
     def make_html_table(rows):
       html = """
@@ -106,7 +292,7 @@ with gr.Blocks() as interface:
                   <thead style='position: sticky; top: 0; background-color: #2c2c2c; z-index: 1;'>
                       <tr>
       """
-      headers = ["Sample ID", "Technique", "Source", "Predicted Location", "Haplogroup", "Inferred Region", "Context Snippet"]
+      headers = ["Sample ID", "Predicted Country", "Country Explanation", "Predicted Sample Type", "Sample Type Explanation", "Sources", "Time cost"]
       html += "".join(
           f"<th style='padding: 10px; border: 1px solid #555; text-align: left; white-space: nowrap;'>{h}</th>"
           for h in headers
@@ -120,11 +306,31 @@ with gr.Blocks() as interface:
               style = "padding: 10px; border: 1px solid #555; vertical-align: top;"
 
               # For specific columns like Haplogroup, force nowrap
-              if header in ["Haplogroup", "Sample ID", "Technique"]:
+              if header in ["Country Explanation", "Sample Type Explanation"]:
+                style += " max-width: 400px; word-wrap: break-word; white-space: normal;"
+              elif header in ["Sample ID", "Predicted Country", "Predicted Sample Type", "Time cost"]:
                   style += " white-space: nowrap; text-overflow: ellipsis; max-width: 200px; overflow: hidden;"
 
-              if header == "Source" and isinstance(col, str) and col.strip().lower().startswith("http"):
-                  col = f"<a href='{col}' target='_blank' style='color: #4ea1f3; text-decoration: underline;'>{col}</a>"
+              # if header == "Sources" and isinstance(col, str) and col.strip().lower().startswith("http"):
+              #     col = f"<a href='{col}' target='_blank' style='color: #4ea1f3; text-decoration: underline;'>{col}</a>"
+
+              #html += f"<td style='{style}'>{col}</td>"
+              if header == "Sources" and isinstance(col, str):
+                  links = [f"<a href='{url.strip()}' target='_blank' style='color: #4ea1f3; text-decoration: underline;'>{url.strip()}</a>" for url in col.strip().split("\n") if url.strip()]
+                  col = "- "+"<br>- ".join(links)
+              elif isinstance(col, str):
+                  # lines = []
+                  # for line in col.split("\n"):
+                  #     line = line.strip()
+                  #     if not line:
+                  #         continue
+                  #     if line.lower().startswith("rag_llm-"):
+                  #         content = line[len("rag_llm-"):].strip()
+                  #         line = f"{content} (Method: RAG_LLM)"
+                  #     lines.append(f"- {line}")
+                  col = col.replace("\n", "<br>")
+                  #col = col.replace("\t", "&nbsp;&nbsp;&nbsp;&nbsp;")
+                  #col = "<br>".join(lines)
 
               html += f"<td style='{style}'>{col}</td>"
           html += "</tr>"
@@ -133,78 +339,174 @@ with gr.Blocks() as interface:
       return html
   
 
+    # def reset_fields():
+    #     global_stop_flag.value = False  # 💡 Add this to reset the flag
+    #     return (
+    #         #gr.update(value=""),  # single_accession
+    #         gr.update(value=""),  # raw_text
+    #         gr.update(value=None), # file_upload
+    #         #gr.update(value=None),  # resume_file
+    #         #gr.update(value="Single Accession"), # inputMode
+    #         gr.update(value=[], visible=True), # output_table
+    #         # gr.update(value="", visible=True), # output_summary
+    #         # gr.update(value="", visible=True), # output_flag
+    #         gr.update(visible=False), # status
+    #         gr.update(visible=False),  # results_group
+    #         gr.update(value="", visible=False),  # usage_display
+    #         gr.update(value="", visible=False),  # progress_box
+    #     )
     def reset_fields():
-        return (
-            gr.update(value=""),  # single_accession
-            gr.update(value=""),  # raw_text
-            gr.update(value=None), # file_upload
-            gr.update(value="Single Accession"), # inputMode
-            gr.update(value=[], visible=True), # output_table
-            gr.update(value="", visible=True), # output_summary
-            gr.update(value="", visible=True), # output_flag
-            gr.update(visible=False), # status
-            gr.update(visible=False)  # results_group
-        )
-
-    inputMode.change(fn=toggle_input_mode, inputs=inputMode, outputs=[single_input_group, batch_input_group])
-    run_button.click(fn=classify_with_loading, inputs=[], outputs=[status])
+      global_stop_flag.value = False  # Reset the stop flag
+
+      return (
+          gr.update(value=""),          # raw_text
+          gr.update(value=None),        # file_upload
+          gr.update(value=[], visible=True),   # output_table
+          gr.update(value="", visible=True),   # status — reset and make visible again
+          gr.update(visible=False),     # results_group
+          gr.update(value="", visible=True),   # usage_display — reset and make visible again
+          gr.update(value="", visible=True),   # progress_box — reset AND visible!
+      )
+    #inputMode.change(fn=toggle_input_mode, inputs=inputMode, outputs=[single_input_group, batch_input_group])
+    #run_button.click(fn=classify_with_loading, inputs=[], outputs=[status])
+    # run_button.click(
+    #     fn=classify_dynamic,
+    #     inputs=[single_accession, file_upload, raw_text, resume_file,user_email,inputMode],
+    #     outputs=[output_table, 
+    #     #output_summary, output_flag, 
+    #     results_group, download_file, usage_display,status, progress_box]
+    # )
+    
+    # run_button.click(
+    #     fn=threaded_batch_runner,
+    #     #inputs=[file_upload, raw_text, resume_file, user_email],
+    #     inputs=[file_upload, raw_text, user_email],
+    #     outputs=[output_table, results_group, download_file, usage_display, status, progress_box]
+    # )
+#     run_button.click(
+#     fn=threaded_batch_runner,
+#     inputs=[file_upload, raw_text, user_email],
+#     outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
+#     every=0.5  # <-- this tells Gradio to expect streaming
+# )
+    # output_table = gr.HTML()
+    # results_group = gr.Group(visible=False)
+    # download_file = gr.File(visible=False)
+    # usage_display = gr.Markdown(visible=False)
+    # status = gr.Markdown(visible=False)
+    # progress_box = gr.Textbox(visible=False)
+
+#     run_button.click(
+#     fn=threaded_batch_runner,
+#     inputs=[file_upload, raw_text, user_email],
+#     outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
+#     every=0.5,  # streaming enabled
+#     show_progress="full"
+# )
+    print("🎯 DEBUG COMPONENT TYPES")
+    print(type(output_table))
+    print(type(results_group))
+    print(type(download_file))
+    print(type(usage_display))
+    print(type(status))
+    print(type(progress_box))
+    
+
+    # interface.stream(
+    #     fn=threaded_batch_runner,
+    #     inputs=[file_upload, raw_text, user_email],
+    #     outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
+    #     trigger=run_button,
+    #     every=0.5,
+    #     show_progress="full",
+    # )
+    interface.queue()  # No arguments here!
+
     run_button.click(
-        fn=classify_dynamic,
-        inputs=[single_accession, file_upload, raw_text, inputMode],
-        outputs=[output_table, output_summary, output_flag, results_group, status]
+        fn=threaded_batch_runner,
+        inputs=[file_upload, raw_text, user_email],
+        outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
+        concurrency_limit=1,  # ✅ correct in Gradio 5.x
+        queue=True,            # ✅ ensure the queue is used
+        #every=0.5
     )
+    
+
+
+
+    stop_button.click(fn=stop_batch, inputs=[], outputs=[status])
+
+    # reset_button.click(
+    #     #fn=reset_fields,
+    #     fn=lambda: (
+    #         gr.update(value=""), gr.update(value=""), gr.update(value=None), gr.update(value=None), gr.update(value="Single Accession"),
+    #         gr.update(value=[], visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(value="", visible=False), gr.update(value="", visible=False)
+    #     ),
+    #     inputs=[],
+    #     outputs=[
+    #         single_accession, raw_text, file_upload, resume_file,inputMode,
+    #         output_table,# output_summary, output_flag,
+    #         status, results_group, usage_display, progress_box
+    #     ]
+    # )
+    #stop_button.click(fn=lambda sf: (gr.update(value="❌ Stopping...", visible=True), setattr(sf, "value", True) or sf), inputs=[gr.State(stop_flag)], outputs=[status, gr.State(stop_flag)])
+
     reset_button.click(
         fn=reset_fields,
         inputs=[],
-        outputs=[
-            single_accession, raw_text, file_upload, inputMode,
-            output_table, output_summary, output_flag,
-            status, results_group
-        ]
-    )
+        #outputs=[raw_text, file_upload, resume_file, output_table, status, results_group, usage_display, progress_box]
+        outputs=[raw_text, file_upload, output_table, status, results_group, usage_display, progress_box]
+    )  
 
     download_button.click(
       fn=mtdna_backend.save_batch_output, 
-      inputs=[output_table, output_summary, output_flag, output_type], 
+      #inputs=[output_table, output_summary, output_flag, output_type], 
+      inputs=[output_table, output_type], 
       outputs=[download_file])
 
+    # submit_feedback.click(
+    #     fn=mtdna_backend.store_feedback_to_google_sheets, 
+    #     inputs=[single_accession, q1, q2, contact], outputs=feedback_status
+    # )
     submit_feedback.click(
-        fn=mtdna_backend.store_feedback_to_google_sheets, inputs=[single_accession, q1, q2, contact], outputs=feedback_status
+        fn=mtdna_backend.store_feedback_to_google_sheets,
+        inputs=[raw_text, q1, q2, contact],
+        outputs=[feedback_status]
     )
-        # Custom CSS styles
-    gr.HTML("""
-    <style>
-      /* Ensures both sections are equally spaced with the same background size */
-      #output-summary, #output-flag {
-          background-color: #f0f4f8; /* Light Grey for both */
-          padding: 20px;
-          border-radius: 10px;
-          margin-top: 10px;
-          width: 100%; /* Ensure full width */
-          min-height: 150px; /* Ensures both have a minimum height */
-          box-sizing: border-box; /* Prevents padding from increasing size */
-          display: flex;
-          flex-direction: column;
-          justify-content: space-between;
-      }
+    #     # Custom CSS styles
+    # gr.HTML("""
+    # <style>
+    #   /* Ensures both sections are equally spaced with the same background size */
+    #   #output-summary, #output-flag {
+    #       background-color: #f0f4f8; /* Light Grey for both */
+    #       padding: 20px;
+    #       border-radius: 10px;
+    #       margin-top: 10px;
+    #       width: 100%; /* Ensure full width */
+    #       min-height: 150px; /* Ensures both have a minimum height */
+    #       box-sizing: border-box; /* Prevents padding from increasing size */
+    #       display: flex;
+    #       flex-direction: column;
+    #       justify-content: space-between;
+    #   }
       
-      /* Specific background colors */
-      #output-summary {
-          background-color: #434a4b; 
-      }
-
-      #output-flag {
-          background-color: #141616; 
-      }
-
-      /* Ensuring they are in a row and evenly spaced */
-      .gradio-row {
-          display: flex;
-          justify-content: space-between;
-          width: 100%;
-      }
-    </style>
-    """)
+    #   /* Specific background colors */
+    #   #output-summary {
+    #       background-color: #434a4b; 
+    #   }
+
+    #   #output-flag {
+    #       background-color: #141616; 
+    #   }
+
+    #   /* Ensuring they are in a row and evenly spaced */
+    #   .gradio-row {
+    #       display: flex;
+    #       justify-content: space-between;
+    #       width: 100%;
+    #   }
+    # </style>
+    # """)
 
 
 interface.launch(share=True,debug=True)

data_preprocess.py

@@ -0,0 +1,625 @@
+import re
+import os
+import streamlit as st
+import subprocess
+import re
+from Bio import Entrez
+from docx import Document
+import fitz
+import spacy
+from spacy.cli import download
+from NER.PDF import pdf
+from NER.WordDoc import wordDoc
+from NER.html import extractHTML
+from NER.word2Vec import word2vec
+from transformers import pipeline
+import urllib.parse, requests
+from pathlib import Path
+import pandas as pd
+from iterate3 import model
+import nltk
+nltk.download('punkt_tab')
+def download_excel_file(url, save_path="temp.xlsx"):
+    if "view.officeapps.live.com" in url:
+        parsed_url = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
+        real_url = urllib.parse.unquote(parsed_url["src"][0])
+        response = requests.get(real_url)
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+        return save_path
+    elif url.startswith("http") and (url.endswith(".xls") or url.endswith(".xlsx")):
+        response = requests.get(url)
+        response.raise_for_status()  # Raises error if download fails
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+            print(len(response.content))
+        return save_path
+    else:
+        print("URL must point directly to an .xls or .xlsx file\n or it already downloaded.")
+        return url
+def extract_text(link,saveFolder):
+  text = ""
+  name = link.split("/")[-1]
+  file_path = Path(saveFolder) / name
+  # pdf
+  if link.endswith(".pdf"):
+    if file_path.is_file():
+      link = saveFolder + "/" + name
+      print("File exists.")
+    p = pdf.PDF(link,saveFolder)
+    text = p.extractTextWithPDFReader()
+    #text_exclude_table = p.extract_text_excluding_tables()
+  # worddoc
+  elif link.endswith(".doc") or link.endswith(".docx"):
+    d = wordDoc.wordDoc(link,saveFolder)
+    text = d.extractTextByPage()
+  # html
+  if link.split(".")[-1].lower() not in "xlsx":
+    if "http" in link or "html" in link:
+      html = extractHTML.HTML("",link)
+      text = html.getListSection() # the text already clean
+  return text
+def extract_table(link,saveFolder):
+  table = []
+  name = link.split("/")[-1]
+  file_path = Path(saveFolder) / name
+  # pdf
+  if link.endswith(".pdf"):
+    if file_path.is_file():
+      link = saveFolder + "/" + name
+      print("File exists.")
+    p = pdf.PDF(link,saveFolder)
+    table = p.extractTable()
+  # worddoc
+  elif link.endswith(".doc") or link.endswith(".docx"):
+    d = wordDoc.wordDoc(link,saveFolder)
+    table = d.extractTableAsList()
+  # excel
+  elif link.split(".")[-1].lower() in "xlsx":
+    # download excel file if it not downloaded yet
+    savePath = saveFolder +"/"+ link.split("/")[-1]
+    excelPath = download_excel_file(link, savePath)
+    try:
+        xls = pd.ExcelFile(excelPath)
+        table_list = []
+        for sheet_name in xls.sheet_names:
+            df = pd.read_excel(xls, sheet_name=sheet_name)
+            cleaned_table = df.fillna("").astype(str).values.tolist()
+            table_list.append(cleaned_table)
+        table = table_list
+    except Exception as e:
+        print("❌ Failed to extract tables from Excel:", e)
+  # html
+  elif "http" in link or "html" in link:
+    html = extractHTML.HTML("",link)
+    table = html.extractTable() # table is a list
+  table = clean_tables_format(table)
+  return table
+
+def clean_tables_format(tables):
+    """
+    Ensures all tables are in consistent format: List[List[List[str]]]
+    Cleans by:
+    - Removing empty strings and rows
+    - Converting all cells to strings
+    - Handling DataFrames and list-of-lists
+    """
+    cleaned = []
+    if tables:
+      for table in tables:
+          standardized = []
+
+          # Case 1: Pandas DataFrame
+          if isinstance(table, pd.DataFrame):
+              table = table.fillna("").astype(str).values.tolist()
+
+          # Case 2: List of Lists
+          if isinstance(table, list) and all(isinstance(row, list) for row in table):
+              for row in table:
+                  filtered_row = [str(cell).strip() for cell in row if str(cell).strip()]
+                  if filtered_row:
+                      standardized.append(filtered_row)
+
+          if standardized:
+              cleaned.append(standardized)
+
+    return cleaned
+
+import json
+import tiktoken  # Optional: for OpenAI token counting
+def normalize_text_for_comparison(s: str) -> str:
+    """
+    Normalizes text for robust comparison by:
+    1. Converting to lowercase.
+    2. Replacing all types of newlines with a single consistent newline (\n).
+    3. Removing extra spaces (e.g., multiple spaces, leading/trailing spaces on lines).
+    4. Stripping leading/trailing whitespace from the entire string.
+    """
+    s = s.lower()
+    s = s.replace('\r\n', '\n') # Handle Windows newlines
+    s = s.replace('\r', '\n')   # Handle Mac classic newlines
+    
+    # Replace sequences of whitespace (including multiple newlines) with a single space
+    # This might be too aggressive if you need to preserve paragraph breaks,
+    # but good for exact word-sequence matching.
+    s = re.sub(r'\s+', ' ', s) 
+    
+    return s.strip()
+def merge_text_and_tables(text, tables, max_tokens=12000, keep_tables=True, tokenizer="cl100k_base", accession_id=None, isolate=None):
+    """
+    Merge cleaned text and table into one string for LLM input.
+    - Avoids duplicating tables already in text
+    - Extracts only relevant rows from large tables
+    - Skips or saves oversized tables
+    """
+    import importlib
+    json = importlib.import_module("json")
+
+    def estimate_tokens(text_str):
+        try:
+            enc = tiktoken.get_encoding(tokenizer)
+            return len(enc.encode(text_str))
+        except:
+            return len(text_str) // 4  # Fallback estimate
+
+    def is_table_relevant(table, keywords, accession_id=None):
+        flat = " ".join(" ".join(row).lower() for row in table)
+        if accession_id and accession_id.lower() in flat:
+            return True    
+        return any(kw.lower() in flat for kw in keywords)
+    preview, preview1 = "",""    
+    llm_input = "## Document Text\n" + text.strip() + "\n"
+    clean_text = normalize_text_for_comparison(text)
+
+    if tables:
+        for idx, table in enumerate(tables):
+          keywords = ["province","district","region","village","location", "country", "region", "origin", "ancient", "modern"]
+          if accession_id:  keywords += [accession_id.lower()]
+          if isolate: keywords += [isolate.lower()]
+          if is_table_relevant(table, keywords, accession_id):
+            if len(table) > 0:
+              for tab in table:
+                preview = " ".join(tab) if tab else ""
+                preview1 = "\n".join(tab) if tab else ""
+                clean_preview = normalize_text_for_comparison(preview)
+                clean_preview1 = normalize_text_for_comparison(preview1)
+                if clean_preview not in clean_text:
+                  if clean_preview1 not in clean_text:
+                    table_str = json.dumps([tab], indent=2)
+                    llm_input += f"## Table {idx+1}\n{table_str}\n"
+    return llm_input.strip()
+
+def preprocess_document(link, saveFolder, accession=None, isolate=None):
+    try:
+      text = extract_text(link, saveFolder)
+    except: text = ""  
+    try: 
+      tables = extract_table(link, saveFolder)
+    except: tables = []  
+    if accession: accession = accession
+    if isolate: isolate = isolate
+    try:
+      final_input = merge_text_and_tables(text, tables, max_tokens=12000, accession_id=accession, isolate=isolate)
+    except: final_input = ""
+    return text, tables, final_input
+
+def extract_sentences(text):
+    sentences = re.split(r'(?<=[.!?])\s+', text)
+    return [s.strip() for s in sentences if s.strip()]
+
+def is_irrelevant_number_sequence(text):
+    if re.search(r'\b[A-Z]{2,}\d+\b|\b[A-Za-z]+\s+\d+\b', text, re.IGNORECASE):
+        return False
+    word_count = len(re.findall(r'\b[A-Za-z]{2,}\b', text))
+    number_count = len(re.findall(r'\b\d[\d\.]*\b', text))
+    total_tokens = len(re.findall(r'\S+', text))
+    if total_tokens > 0 and (word_count / total_tokens < 0.2) and (number_count / total_tokens > 0.5):
+        return True
+    elif re.fullmatch(r'(\d+(\.\d+)?\s*)+', text.strip()):
+        return True
+    return False
+
+def remove_isolated_single_digits(sentence):
+    tokens = sentence.split()
+    filtered_tokens = []
+    for token in tokens:
+        if token == '0' or token == '1':
+            pass
+        else:
+            filtered_tokens.append(token)
+    return ' '.join(filtered_tokens).strip()
+
+def get_contextual_sentences_BFS(text_content, keyword, depth=2):
+    def extract_codes(sentence):
+    # Match codes like 'A1YU101', 'KM1', 'MO6' — at least 2 letters + numbers
+      return [code for code in re.findall(r'\b[A-Z]{2,}[0-9]+\b', sentence, re.IGNORECASE)]
+    sentences = extract_sentences(text_content)
+    relevant_sentences = set()
+    initial_keywords = set()
+
+    # Define a regex to capture codes like A1YU101 or KM1
+    # This pattern looks for an alphanumeric sequence followed by digits at the end of the string
+    code_pattern = re.compile(r'([A-Z0-9]+?)(\d+)$', re.IGNORECASE)
+
+    # Attempt to parse the keyword into its prefix and numerical part using re.search
+    keyword_match = code_pattern.search(keyword)
+
+    keyword_prefix = None
+    keyword_num = None
+
+    if keyword_match:
+        keyword_prefix = keyword_match.group(1).lower()
+        keyword_num = int(keyword_match.group(2))
+
+    for sentence in sentences:
+        sentence_added = False
+
+        # 1. Check for exact match of the keyword
+        if re.search(r'\b' + re.escape(keyword) + r'\b', sentence, re.IGNORECASE):
+            relevant_sentences.add(sentence.strip())
+            initial_keywords.add(keyword.lower())
+            sentence_added = True
+
+        # 2. Check for range patterns (e.g., A1YU101-A1YU137)
+        # The range pattern should be broad enough to capture the full code string within the range.
+        range_matches = re.finditer(r'([A-Z0-9]+-\d+)', sentence, re.IGNORECASE) # More specific range pattern if needed, or rely on full code pattern below
+        range_matches = re.finditer(r'([A-Z0-9]+\d+)-([A-Z0-9]+\d+)', sentence, re.IGNORECASE) # This is the more robust range pattern
+
+        for r_match in range_matches:
+            start_code_str = r_match.group(1)
+            end_code_str = r_match.group(2)
+
+            # CRITICAL FIX: Use code_pattern.search for start_match and end_match
+            start_match = code_pattern.search(start_code_str)
+            end_match = code_pattern.search(end_code_str)
+
+            if keyword_prefix and keyword_num is not None and start_match and end_match:
+                start_prefix = start_match.group(1).lower()
+                end_prefix = end_match.group(1).lower()
+                start_num = int(start_match.group(2))
+                end_num = int(end_match.group(2))
+
+                # Check if the keyword's prefix matches and its number is within the range
+                if keyword_prefix == start_prefix and \
+                   keyword_prefix == end_prefix and \
+                   start_num <= keyword_num <= end_num:
+                    relevant_sentences.add(sentence.strip())
+                    initial_keywords.add(start_code_str.lower())
+                    initial_keywords.add(end_code_str.lower())
+                    sentence_added = True
+                    break # Only need to find one matching range per sentence
+
+        # 3. If the sentence was added due to exact match or range, add all its alphanumeric codes
+        #    to initial_keywords to ensure graph traversal from related terms.
+        if sentence_added:
+          for word in extract_codes(sentence):
+            initial_keywords.add(word.lower())
+
+
+    # Build word_to_sentences mapping for all sentences
+    word_to_sentences = {}
+    for sent in sentences:
+      codes_in_sent = set(extract_codes(sent))
+      for code in codes_in_sent:
+          word_to_sentences.setdefault(code.lower(), set()).add(sent.strip())
+
+
+    # Build the graph
+    graph = {}
+    for sent in sentences:
+      codes = set(extract_codes(sent))
+      for word1 in codes:
+          word1_lower = word1.lower()
+          graph.setdefault(word1_lower, set())
+          for word2 in codes:
+              word2_lower = word2.lower()
+              if word1_lower != word2_lower:
+                  graph[word1_lower].add(word2_lower)
+
+
+    # Perform BFS/graph traversal
+    queue = [(k, 0) for k in initial_keywords if k in word_to_sentences]
+    visited_words = set(initial_keywords)
+
+    while queue:
+        current_word, level = queue.pop(0)
+        if level >= depth:
+            continue
+
+        relevant_sentences.update(word_to_sentences.get(current_word, []))
+
+        for neighbor in graph.get(current_word, []):
+            if neighbor not in visited_words:
+                visited_words.add(neighbor)
+                queue.append((neighbor, level + 1))
+
+    final_sentences = set()
+    for sentence in relevant_sentences:
+        if not is_irrelevant_number_sequence(sentence):
+            processed_sentence = remove_isolated_single_digits(sentence)
+            if processed_sentence:
+                final_sentences.add(processed_sentence)
+
+    return "\n".join(sorted(list(final_sentences)))
+
+
+
+def get_contextual_sentences_DFS(text_content, keyword, depth=2):
+    sentences = extract_sentences(text_content)
+
+    # Build word-to-sentences mapping
+    word_to_sentences = {}
+    for sent in sentences:
+        words_in_sent = set(re.findall(r'\b[A-Za-z0-9\-_\/]+\b', sent))
+        for word in words_in_sent:
+            word_to_sentences.setdefault(word.lower(), set()).add(sent.strip())
+
+    # Function to extract codes in a sentence
+    def extract_codes(sentence):
+      # Only codes like 'KSK1', 'MG272794', not pure numbers
+      return [code for code in re.findall(r'\b[A-Z]{2,}[0-9]+\b', sentence, re.IGNORECASE)]
+
+    # DFS with priority based on distance to keyword and early stop if country found
+    def dfs_traverse(current_word, current_depth, max_depth, visited_words, collected_sentences, parent_sentence=None):
+        country = "unknown"
+        if current_depth > max_depth:
+            return country, False
+
+        if current_word not in word_to_sentences:
+            return country, False
+
+        for sentence in word_to_sentences[current_word]:
+            if sentence == parent_sentence:
+                continue  # avoid reusing the same sentence
+
+            collected_sentences.add(sentence)
+
+            #print("current_word:", current_word)
+            small_sen = extract_context(sentence, current_word, int(len(sentence) / 4))
+            #print(small_sen)
+            country = model.get_country_from_text(small_sen)
+            #print("small context country:", country)
+            if country.lower() != "unknown":
+                return country, True
+            else:
+                country = model.get_country_from_text(sentence)
+                #print("full sentence country:", country)
+                if country.lower() != "unknown":
+                    return country, True
+
+            codes_in_sentence = extract_codes(sentence)
+            idx = next((i for i, code in enumerate(codes_in_sentence) if code.lower() == current_word.lower()), None)
+            if idx is None:
+                continue
+
+            sorted_children = sorted(
+                [code for code in codes_in_sentence if code.lower() not in visited_words],
+                key=lambda x: (abs(codes_in_sentence.index(x) - idx),
+                               0 if codes_in_sentence.index(x) > idx else 1)
+            )
+
+            #print("sorted_children:", sorted_children)
+            for child in sorted_children:
+                child_lower = child.lower()
+                if child_lower not in visited_words:
+                    visited_words.add(child_lower)
+                    country, should_stop = dfs_traverse(
+                        child_lower, current_depth + 1, max_depth,
+                        visited_words, collected_sentences, parent_sentence=sentence
+                    )
+                    if should_stop:
+                        return country, True
+
+        return country, False
+
+    # Begin DFS
+    collected_sentences = set()
+    visited_words = set([keyword.lower()])
+    country, status = dfs_traverse(keyword.lower(), 0, depth, visited_words, collected_sentences)
+
+    # Filter irrelevant sentences
+    final_sentences = set()
+    for sentence in collected_sentences:
+        if not is_irrelevant_number_sequence(sentence):
+            processed = remove_isolated_single_digits(sentence)
+            if processed:
+                final_sentences.add(processed)
+    if not final_sentences:
+      return country, text_content
+    return country, "\n".join(sorted(list(final_sentences)))
+
+# Helper function for normalizing text for overlap comparison
+def normalize_for_overlap(s: str) -> str:
+    s = re.sub(r'[^a-zA-Z0-9\s]', ' ', s).lower()
+    s = re.sub(r'\s+', ' ', s).strip()
+    return s
+
+def merge_texts_skipping_overlap(text1: str, text2: str) -> str:
+    if not text1: return text2
+    if not text2: return text1
+
+    # Case 1: text2 is fully contained in text1 or vice-versa
+    if text2 in text1:
+        return text1
+    if text1 in text2:
+        return text2
+
+    # --- Option 1: Original behavior (suffix of text1, prefix of text2) ---
+    # This is what your function was primarily designed for.
+    # It looks for the overlap at the "junction" of text1 and text2.
+    
+    max_junction_overlap = 0
+    for i in range(min(len(text1), len(text2)), 0, -1):
+        suffix1 = text1[-i:]
+        prefix2 = text2[:i]
+        # Prioritize exact match, then normalized match
+        if suffix1 == prefix2:
+            max_junction_overlap = i
+            break
+        elif normalize_for_overlap(suffix1) == normalize_for_overlap(prefix2):
+            max_junction_overlap = i
+            break # Take the first (longest) normalized match
+
+    if max_junction_overlap > 0:
+        merged_text = text1 + text2[max_junction_overlap:]
+        return re.sub(r'\s+', ' ', merged_text).strip()
+
+    # --- Option 2: Longest Common Prefix (for cases like "Hi, I am Vy.") ---
+    # This addresses your specific test case where the overlap is at the very beginning of both strings.
+    # This is often used when trying to deduplicate content that shares a common start.
+
+    longest_common_prefix_len = 0
+    min_len = min(len(text1), len(text2))
+    for i in range(min_len):
+        if text1[i] == text2[i]:
+            longest_common_prefix_len = i + 1
+        else:
+            break
+    
+    # If a common prefix is found AND it's a significant portion (e.g., more than a few chars)
+    # AND the remaining parts are distinct, then apply this merge.
+    # This is a heuristic and might need fine-tuning.
+    if longest_common_prefix_len > 0 and \
+       text1[longest_common_prefix_len:].strip() and \
+       text2[longest_common_prefix_len:].strip():
+
+        # Only merge this way if the remaining parts are not empty (i.e., not exact duplicates)
+        # For "Hi, I am Vy. Nice to meet you." and "Hi, I am Vy. Goodbye Vy."
+        # common prefix is "Hi, I am Vy."
+        # Remaining text1: " Nice to meet you."
+        # Remaining text2: " Goodbye Vy."
+        # So we merge common_prefix + remaining_text1 + remaining_text2
+        
+        common_prefix_str = text1[:longest_common_prefix_len]
+        remainder_text1 = text1[longest_common_prefix_len:]
+        remainder_text2 = text2[longest_common_prefix_len:]
+        
+        merged_text = common_prefix_str + remainder_text1 + remainder_text2
+        return re.sub(r'\s+', ' ', merged_text).strip()
+
+
+    # If neither specific overlap type is found, just concatenate
+    merged_text = text1 + text2
+    return re.sub(r'\s+', ' ', merged_text).strip()
+
+def save_text_to_docx(text_content: str, file_path: str):
+    """
+    Saves a given text string into a .docx file.
+
+    Args:
+        text_content (str): The text string to save.
+        file_path (str): The full path including the filename where the .docx file will be saved.
+                         Example: '/content/drive/MyDrive/CollectData/Examples/test/SEA_1234/merged_document.docx'
+    """
+    try:
+        document = Document()
+
+        # Add the entire text as a single paragraph, or split by newlines for multiple paragraphs
+        for paragraph_text in text_content.split('\n'):
+            document.add_paragraph(paragraph_text)
+
+        document.save(file_path)
+        print(f"Text successfully saved to '{file_path}'")
+    except Exception as e:
+        print(f"Error saving text to docx file: {e}") 
+
+'''2 scenerios:
+- quick look then found then deepdive and directly get location then stop
+- quick look then found then deepdive but not find location then hold the related words then 
+look another files iteratively for each related word and find location and stop'''
+def extract_context(text, keyword, window=500):
+    # firstly try accession number
+    code_pattern = re.compile(r'([A-Z0-9]+?)(\d+)$', re.IGNORECASE)
+
+    # Attempt to parse the keyword into its prefix and numerical part using re.search
+    keyword_match = code_pattern.search(keyword)
+
+    keyword_prefix = None
+    keyword_num = None
+
+    if keyword_match:
+        keyword_prefix = keyword_match.group(1).lower()
+        keyword_num = int(keyword_match.group(2))
+    text = text.lower()    
+    idx = text.find(keyword.lower())
+    if idx == -1:
+      if keyword_prefix:
+        idx = text.find(keyword_prefix)
+      if idx == -1:
+        return "Sample ID not found."
+      return text[max(0, idx-window): idx+window]  
+    return text[max(0, idx-window): idx+window]
+def process_inputToken(filePaths, saveLinkFolder,accession=None, isolate=None):
+  cache = {}
+  country = "unknown"
+  output = ""
+  tem_output, small_output = "",""
+  keyword_appear = (False,"")
+  keywords = []
+  if isolate: keywords.append(isolate)
+  if accession: keywords.append(accession)
+  for f in filePaths:
+    # scenerio 1: direct location: truncate the context and then use qa model?
+    if keywords:
+      for keyword in keywords:
+        text, tables, final_input = preprocess_document(f,saveLinkFolder, isolate=keyword)
+        if keyword in final_input:
+          context = extract_context(final_input, keyword)
+          # quick look if country already in context and if yes then return
+          country = model.get_country_from_text(context)
+          if country != "unknown":
+            return country, context, final_input
+          else:
+            country = model.get_country_from_text(final_input)  
+            if country != "unknown":
+              return country, context, final_input
+            else: # might be cross-ref
+              keyword_appear = (True, f)
+              cache[f] = context
+              small_output = merge_texts_skipping_overlap(output, context) + "\n"
+              chunkBFS = get_contextual_sentences_BFS(small_output, keyword)
+              countryBFS = model.get_country_from_text(chunkBFS)
+              countryDFS, chunkDFS = get_contextual_sentences_DFS(output, keyword)
+              output = merge_texts_skipping_overlap(output, final_input)
+              if countryDFS != "unknown" and countryBFS != "unknown":
+                if len(chunkDFS) <= len(chunkBFS):
+                  return countryDFS, chunkDFS, output
+                else:
+                  return countryBFS, chunkBFS, output
+              else:        
+                if countryDFS != "unknown":  
+                  return countryDFS, chunkDFS, output
+                if countryBFS != "unknown":
+                  return countryBFS, chunkBFS, output
+        else:
+        # scenerio 2: 
+          '''cross-ref: ex: A1YU101 keyword in file 2 which includes KM1 but KM1 in file 1 
+          but if we look at file 1 first then maybe we can have lookup dict which country 
+          such as Thailand as the key and its re''' 
+          cache[f] = final_input
+          if keyword_appear[0] == True:
+            for c in cache:
+              if c!=keyword_appear[1]:
+                if cache[c].lower() not in output.lower():
+                  output = merge_texts_skipping_overlap(output, cache[c]) + "\n"
+                  chunkBFS = get_contextual_sentences_BFS(output, keyword)
+                  countryBFS = model.get_country_from_text(chunkBFS)
+                  countryDFS, chunkDFS = get_contextual_sentences_DFS(output, keyword)
+                  if countryDFS != "unknown" and countryBFS != "unknown":
+                    if len(chunkDFS) <= len(chunkBFS):
+                      return countryDFS, chunkDFS, output
+                    else:
+                      return countryBFS, chunkBFS, output
+                  else:        
+                    if countryDFS != "unknown":  
+                      return countryDFS, chunkDFS, output
+                    if countryBFS != "unknown":
+                      return countryBFS, chunkBFS, output
+          else:
+            if cache[f].lower() not in output.lower():
+              output = merge_texts_skipping_overlap(output, cache[f]) + "\n"          
+  if len(output) == 0 or keyword_appear[0]==False:
+    for c in cache:
+      if cache[c].lower() not in output.lower():
+        output = merge_texts_skipping_overlap(output, cache[c]) + "\n"
+  return country, "", output               

model.py

@@ -0,0 +1,1255 @@
+import re
+import pycountry
+from docx import Document
+import json
+import os
+import numpy as np
+import faiss
+from collections import defaultdict
+import ast # For literal_eval
+import math # For ceiling function
+from iterate3 import data_preprocess
+import mtdna_classifier
+# --- IMPORTANT: UNCOMMENT AND CONFIGURE YOUR REAL API KEY ---
+import google.generativeai as genai
+os.environ["GOOGLE_API_KEY"] = "AIzaSyDi0CNKBgEtnr6YuPaY6YNEuC5wT0cdKhk"
+genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+
+import nltk
+from nltk.corpus import stopwords
+try:
+    nltk.data.find('corpora/stopwords')
+except LookupError:
+    nltk.download('stopwords')
+nltk.download('punkt_tab')    
+# --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) ---
+# Prices are per 1,000 tokens
+PRICE_PER_1K_INPUT_LLM = 0.000075  # $0.075 per 1M tokens
+PRICE_PER_1K_OUTPUT_LLM = 0.0003   # $0.30 per 1M tokens
+PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
+
+# --- API Functions (REAL API FUNCTIONS) ---
+
+# def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"):
+#     """Generates an embedding for the given text using a Google embedding model."""
+#     try:
+#         result = genai.embed_content(
+#             model="models/text-embedding-004", # Specify the embedding model
+#             content=text,
+#             task_type=task_type
+#         )
+#         return np.array(result['embedding']).astype('float32')
+#     except Exception as e:
+#         print(f"Error getting embedding: {e}")
+#         return np.zeros(768, dtype='float32')
+def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"):
+    """Safe Gemini 1.5 embedding call with fallback."""
+    import numpy as np
+    try:
+        if not text or len(text.strip()) == 0:
+            raise ValueError("Empty text cannot be embedded.")
+        result = genai.embed_content(
+            model="models/text-embedding-004",
+            content=text,
+            task_type=task_type
+        )
+        return np.array(result['embedding'], dtype='float32')
+    except Exception as e:
+        print(f"❌ Embedding error: {e}")
+        return np.zeros(768, dtype='float32')
+
+
+def call_llm_api(prompt, model_name='gemini-1.5-flash-latest'):
+    """Calls a Google Gemini LLM with the given prompt."""
+    try:
+        model = genai.GenerativeModel(model_name)
+        response = model.generate_content(prompt)
+        return response.text, model # Return model instance for token counting
+    except Exception as e:
+        print(f"Error calling LLM: {e}")
+        return "Error: Could not get response from LLM API.", None
+
+
+# --- Core Document Processing Functions (All previously provided and fixed) ---
+
+def read_docx_text(path):
+    """
+    Reads text and extracts potential table-like strings from a .docx document.
+    Separates plain text from structured [ [ ] ] list-like tables.
+    Also attempts to extract a document title.
+    """
+    doc = Document(path)
+    plain_text_paragraphs = []
+    table_strings = []
+    document_title = "Unknown Document Title" # Default
+
+    # Attempt to extract the document title from the first few paragraphs
+    title_paragraphs = [p.text.strip() for p in doc.paragraphs[:5] if p.text.strip()]
+    if title_paragraphs:
+        # A heuristic to find a title: often the first or second non-empty paragraph
+        # or a very long first paragraph if it's the title
+        if len(title_paragraphs[0]) > 50 and "Human Genetics" not in title_paragraphs[0]:
+            document_title = title_paragraphs[0]
+        elif len(title_paragraphs) > 1 and len(title_paragraphs[1]) > 50 and "Human Genetics" not in title_paragraphs[1]:
+            document_title = title_paragraphs[1]
+        elif any("Complete mitochondrial genomes" in p for p in title_paragraphs):
+            # Fallback to a known title phrase if present
+            document_title = "Complete mitochondrial genomes of Thai and Lao populations indicate an ancient origin of Austroasiatic groups and demic diffusion in the spread of Tai–Kadai languages"
+
+    current_table_lines = []
+    in_table_parsing_mode = False
+
+    for p in doc.paragraphs:
+        text = p.text.strip()
+        if not text:
+            continue
+
+        # Condition to start or continue table parsing
+        if text.startswith("## Table "): # Start of a new table section
+            if in_table_parsing_mode and current_table_lines:
+                table_strings.append("\n".join(current_table_lines))
+            current_table_lines = [text] # Include the "## Table X" line
+            in_table_parsing_mode = True
+        elif in_table_parsing_mode and (text.startswith("[") or text.startswith('"')):
+            # Continue collecting lines if we're in table mode and it looks like table data
+            # Table data often starts with '[' for lists, or '"' for quoted strings within lists.
+            current_table_lines.append(text)
+        else:
+            # If not in table mode, or if a line doesn't look like table data,
+            # then close the current table (if any) and add the line to plain text.
+            if in_table_parsing_mode and current_table_lines:
+                table_strings.append("\n".join(current_table_lines))
+                current_table_lines = []
+            in_table_parsing_mode = False
+            plain_text_paragraphs.append(text)
+
+    # After the loop, add any remaining table lines
+    if current_table_lines:
+        table_strings.append("\n".join(current_table_lines))
+
+    return "\n".join(plain_text_paragraphs), table_strings, document_title
+
+# --- Structured Data Extraction and RAG Functions ---
+
+def parse_literal_python_list(table_str):
+    list_match = re.search(r'(\[\s*\[\s*(?:.|\n)*?\s*\]\s*\])', table_str)
+    #print("Debug: list_match object (before if check):", list_match)
+    if not list_match:
+        if "table" in table_str.lower(): # then the table doest have the "]]" at the end
+            table_str += "]]"
+            list_match = re.search(r'(\[\s*\[\s*(?:.|\n)*?\s*\]\s*\])', table_str)
+    if list_match:
+        try:
+            matched_string = list_match.group(1)
+            #print("Debug: Matched string for literal_eval:", matched_string)
+            return ast.literal_eval(matched_string)
+        except (ValueError, SyntaxError) as e:
+            print(f"Error evaluating literal: {e}")
+            return []
+    return []
+
+
+_individual_code_parser = re.compile(r'([A-Z0-9]+?)(\d+)$', re.IGNORECASE)
+def _parse_individual_code_parts(code_str):
+    match = _individual_code_parser.search(code_str)
+    if match:
+        return match.group(1), match.group(2)
+    return None, None
+
+
+def parse_sample_id_to_population_code(plain_text_content):
+    sample_id_map = {}
+    contiguous_ranges_data = defaultdict(list)
+
+    #section_start_marker = "The sample identification of each population is as follows:"
+    section_start_marker = ["The sample identification of each population is as follows:","## table"]
+    
+    for s in section_start_marker:
+      relevant_text_search = re.search(
+          re.escape(s.lower()) + r"\s*(.*?)(?=\n##|\Z)",
+          plain_text_content.lower(),
+          re.DOTALL
+      )
+      if relevant_text_search: 
+        break
+      
+    if not relevant_text_search:
+        print("Warning: 'Sample ID Population Code' section start marker not found or block empty.")
+        return sample_id_map, contiguous_ranges_data
+
+    relevant_text_block = relevant_text_search.group(1).strip()
+
+    # print(f"\nDEBUG_PARSING: --- Start of relevant_text_block (first 500 chars) ---")
+    # print(relevant_text_block[:500])
+    # print(f"DEBUG_PARSING: --- End of relevant_text_block (last 500 chars) ---")
+    # print(relevant_text_block[-500:])
+    # print(f"DEBUG_PARSING: Relevant text block length: {len(relevant_text_block)}")
+
+    mapping_pattern = re.compile(
+    r'\b([A-Z0-9]+\d+)(?:-([A-Z0-9]+\d+))?\s+([A-Z0-9]+)\b', # Changed the last group
+    re.IGNORECASE)
+
+    range_expansion_count = 0
+    direct_id_count = 0
+    total_matches_found = 0
+    for match in mapping_pattern.finditer(relevant_text_block):
+        total_matches_found += 1
+        id1_full_str, id2_full_str_opt, pop_code = match.groups()
+
+        #print(f"  DEBUG_PARSING: Matched: '{match.group(0)}'")
+
+        pop_code_upper = pop_code.upper()
+
+        id1_prefix, id1_num_str = _parse_individual_code_parts(id1_full_str)
+        if id1_prefix is None:
+            #print(f"    DEBUG_PARSING: Failed to parse ID1: {id1_full_str}. Skipping this mapping.")
+            continue
+
+        if id2_full_str_opt:
+            id2_prefix_opt, id2_num_str_opt = _parse_individual_code_parts(id2_full_str_opt)
+            if id2_prefix_opt is None:
+                #print(f"    DEBUG_PARSING: Failed to parse ID2: {id2_full_str_opt}. Treating {id1_full_str} as single ID1.")
+                sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper
+                direct_id_count += 1
+                continue
+
+            #print(f"    DEBUG_PARSING: Comparing prefixes: '{id1_prefix.lower()}' vs '{id2_prefix_opt.lower()}'")
+            if id1_prefix.lower() == id2_prefix_opt.lower():
+                #print(f"    DEBUG_PARSING: ---> Prefixes MATCH for range expansion! Range: {id1_prefix}{id1_num_str}-{id2_prefix_opt}{id2_num_str_opt}")
+                try:
+                    start_num = int(id1_num_str)
+                    end_num = int(id2_num_str_opt)
+                    for num in range(start_num, end_num + 1):
+                        sample_id = f"{id1_prefix.upper()}{num}"
+                        sample_id_map[sample_id] = pop_code_upper
+                        range_expansion_count += 1
+                    contiguous_ranges_data[id1_prefix.upper()].append(
+                        (start_num, end_num, pop_code_upper)
+                    )
+                except ValueError:
+                    print(f"        DEBUG_PARSING: ValueError in range conversion for {id1_num_str}-{id2_num_str_opt}. Adding endpoints only.")
+                    sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper
+                    sample_id_map[f"{id2_prefix_opt.upper()}{id2_num_str_opt}"] = pop_code_upper
+                    direct_id_count += 2
+            else:
+                #print(f"    DEBUG_PARSING: Prefixes MISMATCH for range: '{id1_prefix}' vs '{id2_prefix_opt}'. Adding endpoints only.")
+                sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper
+                sample_id_map[f"{id2_prefix_opt.upper()}{id2_num_str_opt}"] = pop_code_upper
+                direct_id_count += 2
+        else:
+            sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper
+            direct_id_count += 1
+
+    # print(f"DEBUG_PARSING: Total matches found by regex: {total_matches_found}.")
+    # print(f"DEBUG_PARSING: Parsed sample IDs: {len(sample_id_map)} total entries.")
+    # print(f"DEBUG_PARSING:   (including {range_expansion_count} from range expansion and {direct_id_count} direct ID/endpoint entries).")
+    return sample_id_map, contiguous_ranges_data
+
+country_keywords_regional_overrides = {
+    "north thailand": "Thailand", "central thailand": "Thailand",
+    "northeast thailand": "Thailand", "east myanmar": "Myanmar", "west thailand": "Thailand",
+    "central india": "India", "east india": "India", "northeast india": "India",
+    "south sibera": "Russia", "siberia": "Russia", "yunnan": "China", #"tibet": "China",
+    "sumatra": "Indonesia", "borneo": "Indonesia",
+    "northern mindanao": "Philippines", "west malaysia": "Malaysia",
+    "mongolia": "China",
+    "beijing": "China",
+    "north laos": "Laos", "central laos": "Laos",
+    "east myanmar": "Myanmar", "west myanmar": "Myanmar"}
+
+# Updated get_country_from_text function
+def get_country_from_text(text):
+    text_lower = text.lower()
+
+    # 1. Use pycountry for official country names and common aliases
+    for country in pycountry.countries:
+        # Check full name match first
+        if text_lower == country.name.lower():
+            return country.name
+        
+        # Safely check for common_name
+        if hasattr(country, 'common_name') and text_lower == country.common_name.lower():
+            return country.common_name
+            
+        # Safely check for official_name
+        if hasattr(country, 'official_name') and text_lower == country.official_name.lower():
+            return country.official_name
+
+        # Check if country name is part of the text (e.g., 'Thailand' in 'Thailand border')
+        if country.name.lower() in text_lower:
+            return country.name
+            
+        # Safely check if common_name is part of the text
+        if hasattr(country, 'common_name') and country.common_name.lower() in text_lower:
+            return country.common_name
+    # 2. Prioritize specific regional overrides
+    for keyword, country in country_keywords_regional_overrides.items():
+        if keyword in text_lower:
+            return country
+    # 3. Check for broader regions that you want to map to "unknown" or a specific country
+    if "north asia" in text_lower or "southeast asia" in text_lower or "east asia" in text_lower:
+        return "unknown"
+
+    return "unknown"
+
+# Get the list of English stop words from NLTK
+non_meaningful_pop_names = set(stopwords.words('english'))
+
+def parse_population_code_to_country(plain_text_content, table_strings):
+    pop_code_country_map = {}
+    pop_code_ethnicity_map = {} # NEW: To store ethnicity for structured lookup
+    pop_code_specific_loc_map = {} # NEW: To store specific location for structured lookup
+
+    # Regex for parsing population info in structured lists and general text
+    # This pattern captures: (Pop Name/Ethnicity) (Pop Code) (Region/Specific Location) (Country) (Linguistic Family)
+    # The 'Pop Name/Ethnicity' (Group 1) is often the ethnicity
+    pop_info_pattern = re.compile(
+          r'([A-Za-z\s]+?)\s+([A-Z]+\d*)\s+'      # Pop Name (Group 1), Pop Code (Group 2) - Changed \d+ to \d* for codes like 'SH'
+          r'([A-Za-z\s\(\)\-,\/]+?)\s+'          # Region/Specific Location (Group 3)
+          r'(North+|South+|West+|East+|Thailand|Laos|Cambodia|Myanmar|Philippines|Indonesia|Malaysia|China|India|Taiwan|Vietnam|Russia|Nepal|Japan|South Korea)\b' # Country (Group 4)
+          r'(?:.*?([A-Za-z\s\-]+))?\s*'          # Optional Linguistic Family (Group 5), made optional with ?, followed by optional space
+          r'(\d+(?:\s+\d+\.?\d*)*)?', # Match all the numbers (Group 6) - made optional
+          re.IGNORECASE
+      )
+    for table_str in table_strings:
+        table_data = parse_literal_python_list(table_str)
+        if table_data:
+            is_list_of_lists = bool(table_data) and isinstance(table_data[0], list)
+            if is_list_of_lists:
+                for row_idx, row in enumerate(table_data):
+                    row_text = " ".join(map(str, row))
+                    match = pop_info_pattern.search(row_text)
+                    if match:
+                        pop_name = match.group(1).strip()
+                        pop_code = match.group(2).upper()
+                        specific_loc_text = match.group(3).strip()
+                        country_text = match.group(4).strip()
+                        linguistic_family = match.group(5).strip() if match.group(5) else 'unknown'
+
+                        final_country = get_country_from_text(country_text)
+                        if final_country == 'unknown': # Try specific loc text for country if direct country is not found
+                            final_country = get_country_from_text(specific_loc_text)
+
+                        if pop_code:
+                            pop_code_country_map[pop_code] = final_country
+
+                            # Populate ethnicity map (often Pop Name is ethnicity)
+                            pop_code_ethnicity_map[pop_code] = pop_name
+
+                            # Populate specific location map
+                            pop_code_specific_loc_map[pop_code] = specific_loc_text # Store as is from text
+            else:
+                row_text = " ".join(map(str, table_data))   
+                match = pop_info_pattern.search(row_text)
+                if match:
+                    pop_name = match.group(1).strip()
+                    pop_code = match.group(2).upper()
+                    specific_loc_text = match.group(3).strip()
+                    country_text = match.group(4).strip()
+                    linguistic_family = match.group(5).strip() if match.group(5) else 'unknown'
+
+                    final_country = get_country_from_text(country_text)
+                    if final_country == 'unknown': # Try specific loc text for country if direct country is not found
+                        final_country = get_country_from_text(specific_loc_text)
+
+                    if pop_code:
+                        pop_code_country_map[pop_code] = final_country
+
+                        # Populate ethnicity map (often Pop Name is ethnicity)
+                        pop_code_ethnicity_map[pop_code] = pop_name
+
+                        # Populate specific location map
+                        pop_code_specific_loc_map[pop_code] = specific_loc_text # Store as is from text
+
+                        # # Special case refinements for ethnicity/location if more specific rules are known from document:
+                        # if pop_name.lower() == "khon mueang": # and specific conditions if needed
+                        #     pop_code_ethnicity_map[pop_code] = "Khon Mueang"
+                        #     # If Khon Mueang has a specific city/district, add here
+                        #     # e.g., if 'Chiang Mai' is directly linked to KM1 in a specific table
+                        #     # pop_code_specific_loc_map[pop_code] = "Chiang Mai"
+                        # elif pop_name.lower() == "lawa":
+                        #      pop_code_ethnicity_map[pop_code] = "Lawa"
+                        # # Add similar specific rules for other populations (e.g., Mon for MO1, MO2, MO3)
+                        # elif pop_name.lower() == "mon":
+                        #     pop_code_ethnicity_map[pop_code] = "Mon"
+                        #     # For MO2: "West Thailand (Thailand Myanmar border)" -> no city
+                        #     # For MO3: "East Myanmar (Thailand Myanmar border)" -> no city
+                        #     # If the doc gives "Bangkok" for MO4, add it here for MO4's actual specific_location.
+                        # # etc.
+
+    # Fallback to parsing general plain text content (sentences)
+    sentences = data_preprocess.extract_sentences(plain_text_content)
+    for s in sentences: # Still focusing on just this one sentence
+      # Use re.finditer to get all matches
+      matches = pop_info_pattern.finditer(s)
+      pop_name, pop_code, specific_loc_text, country_text = "unknown", "unknown", "unknown", "unknown"
+      for match in matches:
+          if match.group(1):
+            pop_name = match.group(1).strip()
+          if match.group(2):  
+            pop_code = match.group(2).upper()
+          if match.group(3):  
+            specific_loc_text = match.group(3).strip()
+          if match.group(4):  
+            country_text = match.group(4).strip()
+          # linguistic_family = match.group(5).strip() if match.group(5) else 'unknown' # Already captured by pop_info_pattern
+
+          final_country = get_country_from_text(country_text)
+          if final_country == 'unknown':
+              final_country = get_country_from_text(specific_loc_text)
+
+          if pop_code.lower() not in non_meaningful_pop_names:
+            if final_country.lower() not in non_meaningful_pop_names:
+              pop_code_country_map[pop_code] = final_country
+            if pop_name.lower() not in non_meaningful_pop_names:  
+              pop_code_ethnicity_map[pop_code] = pop_name # Default ethnicity from Pop Name
+            if specific_loc_text.lower() not in non_meaningful_pop_names:  
+              pop_code_specific_loc_map[pop_code] = specific_loc_text
+
+              # Specific rules for ethnicity/location in plain text:
+              if pop_name.lower() == "khon mueang":
+                  pop_code_ethnicity_map[pop_code] = "Khon Mueang"
+              elif pop_name.lower() == "lawa":
+                  pop_code_ethnicity_map[pop_code] = "Lawa"
+              elif pop_name.lower() == "mon":
+                  pop_code_ethnicity_map[pop_code] = "Mon"
+              elif pop_name.lower() == "seak": # Added specific rule for Seak
+                  pop_code_ethnicity_map[pop_code] = "Seak"
+              elif pop_name.lower() == "nyaw": # Added specific rule for Nyaw
+                  pop_code_ethnicity_map[pop_code] = "Nyaw"
+              elif pop_name.lower() == "nyahkur": # Added specific rule for Nyahkur
+                  pop_code_ethnicity_map[pop_code] = "Nyahkur"
+              elif pop_name.lower() == "suay": # Added specific rule for Suay
+                  pop_code_ethnicity_map[pop_code] = "Suay"
+              elif pop_name.lower() == "soa": # Added specific rule for Soa
+                  pop_code_ethnicity_map[pop_code] = "Soa"
+              elif pop_name.lower() == "bru": # Added specific rule for Bru
+                  pop_code_ethnicity_map[pop_code] = "Bru"
+              elif pop_name.lower() == "khamu": # Added specific rule for Khamu
+                  pop_code_ethnicity_map[pop_code] = "Khamu"
+
+    return pop_code_country_map, pop_code_ethnicity_map, pop_code_specific_loc_map
+
+def general_parse_population_code_to_country(plain_text_content, table_strings):
+    pop_code_country_map = {}
+    pop_code_ethnicity_map = {}
+    pop_code_specific_loc_map = {}
+    sample_id_to_pop_code = {}
+
+    for table_str in table_strings:
+        table_data = parse_literal_python_list(table_str)
+        if not table_data or not isinstance(table_data[0], list):
+            continue
+
+        header_row = [col.lower() for col in table_data[0]]
+        header_map = {col: idx for idx, col in enumerate(header_row)}
+
+        # MJ17: Direct PopCode → Country
+        if 'id' in header_map and 'country' in header_map:
+            for row in table_strings[1:]:
+                row = parse_literal_python_list(row)[0]
+                if len(row) < len(header_row):
+                    continue
+                pop_code = str(row[header_map['id']]).strip()
+                country = str(row[header_map['country']]).strip()
+                province = row[header_map['province']].strip() if 'province' in header_map else 'unknown'
+                pop_group = row[header_map['population group / region']].strip() if 'population group / region' in header_map else 'unknown'
+                pop_code_country_map[pop_code] = country
+                pop_code_specific_loc_map[pop_code] = province
+                pop_code_ethnicity_map[pop_code] = pop_group
+
+        # A1YU101 or EBK/KSK: SampleID → PopCode
+        elif 'sample id' in header_map and 'population code' in header_map:
+            for row in table_strings[1:]:
+                row = parse_literal_python_list(row)[0]
+                if len(row) < 2:
+                    continue
+                sample_id = row[header_map['sample id']].strip().upper()
+                pop_code = row[header_map['population code']].strip().upper()
+                sample_id_to_pop_code[sample_id] = pop_code
+
+        # PopCode → Country (A1YU101/EBK mapping)
+        elif 'population code' in header_map and 'country' in header_map:
+            for row in table_strings[1:]:
+                row = parse_literal_python_list(row)[0]
+                if len(row) < 2:
+                    continue
+                pop_code = row[header_map['population code']].strip().upper()
+                country = row[header_map['country']].strip()
+                pop_code_country_map[pop_code] = country
+
+    return pop_code_country_map, pop_code_ethnicity_map, pop_code_specific_loc_map, sample_id_to_pop_code
+
+def chunk_text(text, chunk_size=500, overlap=50):
+    """Splits text into chunks (by words) with overlap."""
+    chunks = []
+    words = text.split()
+    num_words = len(words)
+
+    start = 0
+    while start < num_words:
+        end = min(start + chunk_size, num_words)
+        chunk = " ".join(words[start:end])
+        chunks.append(chunk)
+
+        if end == num_words:
+            break
+        start += chunk_size - overlap # Move start by (chunk_size - overlap)
+    return chunks
+
+def build_vector_index_and_data(doc_path, index_path="faiss_index.bin", chunks_path="document_chunks.json", structured_path="structured_lookup.json"):
+    """
+    Reads document, builds structured lookup, chunks remaining text, embeds chunks,
+    and builds/saves a FAISS index.
+    """
+    print("Step 1: Reading document and extracting structured data...")
+    # plain_text_content, table_strings, document_title = read_docx_text(doc_path) # Get document_title here
+
+    # sample_id_map, contiguous_ranges_data = parse_sample_id_to_population_code(plain_text_content)
+    # pop_code_to_country, pop_code_to_ethnicity, pop_code_to_specific_loc = parse_population_code_to_country(plain_text_content, table_strings)
+
+    # master_structured_lookup = {}
+    # master_structured_lookup['document_title'] = document_title # Store document title
+    # master_structured_lookup['sample_id_map'] = sample_id_map
+    # master_structured_lookup['contiguous_ranges'] = dict(contiguous_ranges_data)
+    # master_structured_lookup['pop_code_to_country'] = pop_code_to_country
+    # master_structured_lookup['pop_code_to_ethnicity'] = pop_code_to_ethnicity # NEW: Store pop_code to ethnicity map
+    # master_structured_lookup['pop_code_to_specific_loc'] = pop_code_to_specific_loc # NEW: Store pop_code to specific_loc map
+
+
+    # # Final consolidation: Use sample_id_map to derive full info for queries
+    # final_structured_entries = {}
+    # for sample_id, pop_code in master_structured_lookup['sample_id_map'].items():
+    #     country = master_structured_lookup['pop_code_to_country'].get(pop_code, 'unknown')
+    #     ethnicity = master_structured_lookup['pop_code_to_ethnicity'].get(pop_code, 'unknown') # Retrieve ethnicity
+    #     specific_location = master_structured_lookup['pop_code_to_specific_loc'].get(pop_code, 'unknown') # Retrieve specific location
+
+    #     final_structured_entries[sample_id] = {
+    #         'population_code': pop_code,
+    #         'country': country,
+    #         'type': 'modern',
+    #         'ethnicity': ethnicity, # Store ethnicity
+    #         'specific_location': specific_location # Store specific location
+    #     }
+    # master_structured_lookup['final_structured_entries'] = final_structured_entries
+    plain_text_content, table_strings, document_title = read_docx_text(doc_path)
+    pop_code_to_country, pop_code_to_ethnicity, pop_code_to_specific_loc, sample_id_map = general_parse_population_code_to_country(plain_text_content, table_strings)
+
+    final_structured_entries = {}
+    if sample_id_map:
+        for sample_id, pop_code in sample_id_map.items():
+            country = pop_code_to_country.get(pop_code, 'unknown')
+            ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown')
+            specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown')
+            final_structured_entries[sample_id] = {
+                'population_code': pop_code,
+                'country': country,
+                'type': 'modern',
+                'ethnicity': ethnicity,
+                'specific_location': specific_loc
+            }
+    else:
+        for pop_code in pop_code_to_country.keys():
+            country = pop_code_to_country.get(pop_code, 'unknown')
+            ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown')
+            specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown')
+            final_structured_entries[pop_code] = {
+                'population_code': pop_code,
+                'country': country,
+                'type': 'modern',
+                'ethnicity': ethnicity,
+                'specific_location': specific_loc
+            }
+    if not final_structured_entries:
+      # traditional way of A1YU101
+      sample_id_map, contiguous_ranges_data = parse_sample_id_to_population_code(plain_text_content)
+      pop_code_to_country, pop_code_to_ethnicity, pop_code_to_specific_loc = parse_population_code_to_country(plain_text_content, table_strings)
+      if sample_id_map:
+        for sample_id, pop_code in sample_id_map.items():
+            country = pop_code_to_country.get(pop_code, 'unknown')
+            ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown')
+            specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown')
+            final_structured_entries[sample_id] = {
+                'population_code': pop_code,
+                'country': country,
+                'type': 'modern',
+                'ethnicity': ethnicity,
+                'specific_location': specific_loc
+            }
+      else:
+          for pop_code in pop_code_to_country.keys():
+              country = pop_code_to_country.get(pop_code, 'unknown')
+              ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown')
+              specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown')
+              final_structured_entries[pop_code] = {
+                  'population_code': pop_code,
+                  'country': country,
+                  'type': 'modern',
+                  'ethnicity': ethnicity,
+                  'specific_location': specific_loc
+              }
+    
+    master_lookup = {
+        'document_title': document_title,
+        'pop_code_to_country': pop_code_to_country,
+        'pop_code_to_ethnicity': pop_code_to_ethnicity,
+        'pop_code_to_specific_loc': pop_code_to_specific_loc,
+        'sample_id_map': sample_id_map,
+        'final_structured_entries': final_structured_entries
+    }
+    print(f"Structured lookup built with {len(final_structured_entries)} entries in 'final_structured_entries'.")
+
+    with open(structured_path, 'w') as f:
+        json.dump(master_lookup, f, indent=4)
+    print(f"Structured lookup saved to {structured_path}.")
+
+    print("Step 2: Chunking document for RAG vector index...")
+    # replace the chunk here with the all_output from process_inputToken and fallback to this traditional chunk
+    clean_text, clean_table = "", ""
+    if plain_text_content:
+      clean_text = data_preprocess.normalize_for_overlap(plain_text_content)
+    if table_strings:
+      clean_table = data_preprocess.normalize_for_overlap(". ".join(table_strings))
+    all_clean_chunk = clean_text + clean_table
+    document_chunks = chunk_text(all_clean_chunk)
+    print(f"Document chunked into {len(document_chunks)} chunks.")
+    
+    print("Step 3: Generating embeddings for chunks (this might take time and cost API calls)...")
+
+    embedding_model_for_chunks = genai.GenerativeModel('models/text-embedding-004')
+
+    chunk_embeddings = []
+    for i, chunk in enumerate(document_chunks):
+        embedding = get_embedding(chunk, task_type="RETRIEVAL_DOCUMENT")
+        if embedding is not None and embedding.shape[0] > 0:
+            chunk_embeddings.append(embedding)
+        else:
+            print(f"Warning: Failed to get valid embedding for chunk {i}. Skipping.")
+            chunk_embeddings.append(np.zeros(768, dtype='float32'))
+
+    if not chunk_embeddings:
+        raise ValueError("No valid embeddings generated. Check get_embedding function and API.")
+
+    embedding_dimension = chunk_embeddings[0].shape[0]
+    index = faiss.IndexFlatL2(embedding_dimension)
+    index.add(np.array(chunk_embeddings))
+
+    faiss.write_index(index, index_path)
+    with open(chunks_path, "w") as f:
+        json.dump(document_chunks, f)
+
+    print(f"FAISS index built and saved to {index_path}.")
+    print(f"Document chunks saved to {chunks_path}.")
+    return master_lookup, index, document_chunks, all_clean_chunk
+
+
+def load_rag_assets(index_path="faiss_index.bin", chunks_path="document_chunks.json", structured_path="structured_lookup.json"):
+    """Loads pre-built RAG assets (FAISS index, chunks, structured lookup)."""
+    print("Loading RAG assets...")
+    master_structured_lookup = {}
+    if os.path.exists(structured_path):
+        with open(structured_path, 'r') as f:
+            master_structured_lookup = json.load(f)
+        print("Structured lookup loaded.")
+    else:
+        print("Structured lookup file not found. Rebuilding is likely needed.")
+
+    index = None
+    chunks = []
+    if os.path.exists(index_path) and os.path.exists(chunks_path):
+        try:
+            index = faiss.read_index(index_path)
+            with open(chunks_path, "r") as f:
+                chunks = json.load(f)
+            print("FAISS index and chunks loaded.")
+        except Exception as e:
+            print(f"Error loading FAISS index or chunks: {e}. Will rebuild.")
+            index = None
+            chunks = []
+    else:
+        print("FAISS index or chunks files not found.")
+
+    return master_structured_lookup, index, chunks
+# Helper function for query_document_info
+def exactInContext(text, keyword):
+# try keyword_prfix
+  # code_pattern = re.compile(r'([A-Z0-9]+?)(\d+)$', re.IGNORECASE)
+  # # Attempt to parse the keyword into its prefix and numerical part using re.search
+  # keyword_match = code_pattern.search(keyword)
+  # keyword_prefix = None
+  # keyword_num = None
+  # if keyword_match:
+  #     keyword_prefix = keyword_match.group(1).lower()
+  #     keyword_num = int(keyword_match.group(2))
+  text = text.lower()
+  idx = text.find(keyword.lower())
+  if idx == -1:
+    # if keyword_prefix:
+    #   idx = text.find(keyword_prefix)
+    # if idx == -1:
+    #   return False
+    return False
+  return True
+def chooseContextLLM(contexts, kw):
+  # if kw in context
+  for con in contexts:
+    context = contexts[con]
+    if context:
+      if exactInContext(context, kw):
+        return con, context    
+  #if cannot find anything related to kw in context, return all output
+  if contexts["all_output"]:
+    return "all_output", contexts["all_output"]
+  else:
+    # if all_output not exist
+    # look of chunk and still not exist return document chunk
+    if contexts["chunk"]: return "chunk", contexts["chunk"]
+    elif contexts["document_chunk"]:  return "document_chunk", contexts["document_chunk"]
+    else: return None, None  
+def clean_llm_output(llm_response_text, output_format_str):
+    results = []
+    lines = llm_response_text.strip().split('\n')
+    output_country, output_type, output_ethnicity, output_specific_location = [],[],[],[]
+    for line in lines:
+        extracted_country, extracted_type, extracted_ethnicity, extracted_specific_location = "unknown", "unknown", "unknown", "unknown"
+        line = line.strip()
+        if output_format_str == "ethnicity, specific_location/unknown": # Targeted RAG output
+            parsed_output = re.search(r'^\s*([^,]+?),\s*(.+?)\s*$', llm_response_text)
+            if parsed_output:
+                extracted_ethnicity = parsed_output.group(1).strip()
+                extracted_specific_location = parsed_output.group(2).strip()
+            else:
+                print("  DEBUG: LLM did not follow expected 2-field format for targeted RAG. Defaulting to unknown for ethnicity/specific_location.")
+                extracted_ethnicity = 'unknown'
+                extracted_specific_location = 'unknown'
+        elif output_format_str == "modern/ancient/unknown, ethnicity, specific_location/unknown":
+          parsed_output = re.search(r'^\s*([^,]+?),\s*([^,]+?),\s*(.+?)\s*$', llm_response_text)
+          if parsed_output:
+              extracted_type = parsed_output.group(1).strip()
+              extracted_ethnicity = parsed_output.group(2).strip()
+              extracted_specific_location = parsed_output.group(3).strip()
+          else:
+              # Fallback: check if only 2 fields
+              parsed_output_2_fields = re.search(r'^\s*([^,]+?),\s*([^,]+?)\s*$', llm_response_text)
+              if parsed_output_2_fields:
+                  extracted_type = parsed_output_2_fields.group(1).strip()
+                  extracted_ethnicity = parsed_output_2_fields.group(2).strip()
+                  extracted_specific_location = 'unknown'
+              else:
+                  # even simpler fallback: 1 field only
+                  parsed_output_1_field = re.search(r'^\s*([^,]+?)\s*$', llm_response_text)
+                  if parsed_output_1_field:
+                      extracted_type = parsed_output_1_field.group(1).strip()
+                      extracted_ethnicity = 'unknown'
+                      extracted_specific_location = 'unknown'
+                  else:
+                      print("  DEBUG: LLM did not follow any expected simplified format. Attempting verbose parsing fallback.")
+                      type_match_fallback = re.search(r'Type:\s*([A-Za-z\s-]+)', llm_response_text)
+                      extracted_type = type_match_fallback.group(1).strip() if type_match_fallback else 'unknown'
+                      extracted_ethnicity = 'unknown'
+                      extracted_specific_location = 'unknown'
+        else:
+          parsed_output = re.search(r'^\s*([^,]+?),\s*([^,]+?),\s*([^,]+?),\s*(.+?)\s*$', line)
+          if parsed_output:
+              extracted_country = parsed_output.group(1).strip()
+              extracted_type = parsed_output.group(2).strip()
+              extracted_ethnicity = parsed_output.group(3).strip()
+              extracted_specific_location = parsed_output.group(4).strip()
+          else:
+              print(f"  DEBUG: Line did not follow expected 4-field format: {line}")
+              parsed_output_2_fields = re.search(r'^\s*([^,]+?),\s*([^,]+?)\s*$', line)
+              if parsed_output_2_fields:
+                  extracted_country = parsed_output_2_fields.group(1).strip()
+                  extracted_type = parsed_output_2_fields.group(2).strip()
+                  extracted_ethnicity = 'unknown'
+                  extracted_specific_location = 'unknown'
+              else:
+                  print(f"  DEBUG: Fallback to verbose-style parsing: {line}")
+                  country_match_fallback = re.search(r'Country:\s*([A-Za-z\s-]+)', line)
+                  type_match_fallback = re.search(r'Type:\s*([A-Za-z\s-]+)', line)
+                  extracted_country = country_match_fallback.group(1).strip() if country_match_fallback else 'unknown'
+                  extracted_type = type_match_fallback.group(1).strip() if type_match_fallback else 'unknown'
+                  extracted_ethnicity = 'unknown'
+                  extracted_specific_location = 'unknown'
+
+        results.append({
+            "country": extracted_country,
+            "type": extracted_type,
+            "ethnicity": extracted_ethnicity,
+            "specific_location": extracted_specific_location
+            #"country_explain":extracted_country_explain,
+            #"type_explain": extracted_type_explain
+        })
+    # if more than 2 results
+    if output_format_str == "ethnicity, specific_location/unknown":
+      for result in results:
+        if result["ethnicity"] not in output_ethnicity:
+          output_ethnicity.append(result["ethnicity"])
+        if result["specific_location"] not in output_specific_location:  
+          output_specific_location.append(result["specific_location"])
+      return " or ".join(output_ethnicity), " or ".join(output_specific_location)     
+    elif output_format_str == "modern/ancient/unknown, ethnicity, specific_location/unknown":
+      for result in results:
+        if result["type"] not in output_type:
+          output_type.append(result["type"])
+        if result["ethnicity"] not in output_ethnicity:
+          output_ethnicity.append(result["ethnicity"])
+        if result["specific_location"] not in output_specific_location:  
+          output_specific_location.append(result["specific_location"])
+
+      return " or ".join(output_type)," or ".join(output_ethnicity), " or ".join(output_specific_location)    
+    else:
+      for result in results:
+        if result["country"] not in output_country:
+          output_country.append(result["country"])
+        if result["type"] not in output_type:
+          output_type.append(result["type"])
+        if result["ethnicity"] not in output_ethnicity:
+          output_ethnicity.append(result["ethnicity"])
+        if result["specific_location"] not in output_specific_location:  
+          output_specific_location.append(result["specific_location"])
+      return " or ".join(output_country)," or ".join(output_type)," or ".join(output_ethnicity), " or ".join(output_specific_location)           
+
+def parse_multi_sample_llm_output(raw_response: str, output_format_str):
+    """
+    Parse LLM output with possibly multiple metadata lines + shared explanations.
+    """
+    lines = [line.strip() for line in raw_response.strip().splitlines() if line.strip()]
+    metadata_list = []
+    explanation_lines = []
+    if output_format_str == "country_name, modern/ancient/unknown":
+        parts = [x.strip() for x in lines[0].split(",")]
+        if len(parts)==2:
+          metadata_list.append({
+              "country": parts[0],
+              "sample_type": parts[1]#,
+              #"ethnicity": parts[2],
+              #"location": parts[3]
+          })
+        if 1<len(lines):
+          line = lines[1]
+          if "\n" in line:  line = line.split("\n")
+          if ". " in line: line = line.split(". ")
+          if isinstance(line,str): line = [line]
+          explanation_lines += line
+    elif output_format_str == "modern/ancient/unknown":
+      metadata_list.append({
+          "country": "unknown",
+          "sample_type": lines[0]#,
+          #"ethnicity": parts[2],
+          #"location": parts[3]
+      })
+      explanation_lines.append(lines[1])
+
+    # Assign explanations (optional) to each sample — same explanation reused
+    for md in metadata_list:
+        md["country_explanation"] = None
+        md["sample_type_explanation"] = None
+
+        if md["country"].lower() != "unknown" and len(explanation_lines) >= 1:
+            md["country_explanation"] = explanation_lines[0]
+
+        if md["sample_type"].lower() != "unknown":
+            if len(explanation_lines) >= 2:
+                md["sample_type_explanation"] = explanation_lines[1]
+            elif len(explanation_lines) == 1 and md["country"].lower() == "unknown":
+                md["sample_type_explanation"] = explanation_lines[0]
+            elif len(explanation_lines) == 1:
+                md["sample_type_explanation"] = explanation_lines[0]
+    return metadata_list
+
+def merge_metadata_outputs(metadata_list):
+    """
+    Merge a list of metadata dicts into one, combining differing values with 'or'.
+    Assumes all dicts have the same keys.
+    """
+    if not metadata_list:
+        return {}
+
+    merged = {}
+    keys = metadata_list[0].keys()
+
+    for key in keys:
+        values = [md[key] for md in metadata_list if key in md]
+        unique_values = list(dict.fromkeys(values))  # preserve order, remove dupes
+        if "unknown" in unique_values:
+          unique_values.pop(unique_values.index("unknown"))
+        if len(unique_values) == 1:
+            merged[key] = unique_values[0]
+        else:
+            merged[key] = " or ".join(unique_values)
+
+    return merged
+
+
+def query_document_info(query_word, alternative_query_word, metadata, master_structured_lookup, faiss_index, document_chunks, llm_api_function, chunk=None, all_output=None):
+    """
+    Queries the document using a hybrid approach:
+    1. Local structured lookup (fast, cheap, accurate for known patterns).
+    2. RAG with semantic search and LLM (general, flexible, cost-optimized).
+    """
+    if metadata:
+      extracted_country, extracted_specific_location, extracted_ethnicity, extracted_type = metadata["country"], metadata["specific_location"], metadata["ethnicity"], metadata["sample_type"]
+      extracted_col_date, extracted_iso, extracted_title, extracted_features = metadata["collection_date"], metadata["isolate"], metadata["title"], metadata["all_features"]
+    else:
+      extracted_country, extracted_specific_location, extracted_ethnicity, extracted_type = "unknown", "unknown", "unknown", "unknown"
+      extracted_col_date, extracted_iso, extracted_title = "unknown", "unknown", "unknown"
+    # --- NEW: Pre-process alternative_query_word to remove '.X' suffix if present ---
+    if alternative_query_word:
+        alternative_query_word_cleaned = alternative_query_word.split('.')[0]
+    else:
+        alternative_query_word_cleaned = alternative_query_word
+    country_explanation, sample_type_explanation = None, None
+
+    # Use the consolidated final_structured_entries for direct lookup
+    final_structured_entries = master_structured_lookup.get('final_structured_entries', {})
+    document_title = master_structured_lookup.get('document_title', 'Unknown Document Title') # Retrieve document title
+
+    # Default values for all extracted fields. These will be updated.
+    method_used = 'unknown' # Will be updated based on the method that yields a result
+    population_code_from_sl = 'unknown' # To pass to RAG prompt if available
+    total_query_cost = 0
+    # Attempt 1: Try primary query_word (e.g., isolate name) with structured lookup
+    structured_info = final_structured_entries.get(query_word.upper())
+    if structured_info:
+        if extracted_country == 'unknown':
+          extracted_country = structured_info['country']
+        if extracted_type == 'unknown':
+          extracted_type = structured_info['type']
+          
+        # if extracted_ethnicity == 'unknown':
+        #   extracted_ethnicity = structured_info.get('ethnicity', 'unknown') # Get ethnicity from structured lookup
+        # if extracted_specific_location == 'unknown':
+        #   extracted_specific_location = structured_info.get('specific_location', 'unknown') # Get specific_location from structured lookup
+        population_code_from_sl = structured_info['population_code']
+        method_used = "structured_lookup_direct"
+        print(f"'{query_word}' found in structured lookup (direct match).")
+
+    # Attempt 2: Try primary query_word with heuristic range lookup if direct fails (only if not already resolved)
+    if method_used == 'unknown':
+        query_prefix, query_num_str = _parse_individual_code_parts(query_word)
+        if query_prefix is not None and query_num_str is not None:
+            try: query_num = int(query_num_str)
+            except ValueError: query_num = None
+            if query_num is not None:
+                query_prefix_upper = query_prefix.upper()
+                contiguous_ranges = master_structured_lookup.get('contiguous_ranges', defaultdict(list))
+                pop_code_to_country = master_structured_lookup.get('pop_code_to_country', {})
+                pop_code_to_ethnicity = master_structured_lookup.get('pop_code_to_ethnicity', {})
+                pop_code_to_specific_loc = master_structured_lookup.get('pop_code_to_specific_loc', {})
+
+                if query_prefix_upper in contiguous_ranges:
+                    for start_num, end_num, pop_code_for_range in contiguous_ranges[query_prefix_upper]:
+                        if start_num <= query_num <= end_num:
+                            country_from_heuristic = pop_code_to_country.get(pop_code_for_range, 'unknown')
+                            if country_from_heuristic != 'unknown':
+                                if extracted_country == 'unknown':
+                                  extracted_country = country_from_heuristic
+                                if extracted_type == 'unknown':
+                                  extracted_type = 'modern'
+                                # if extracted_ethnicity == 'unknown':
+                                #   extracted_ethnicity = pop_code_to_ethnicity.get(pop_code_for_range, 'unknown')
+                                # if extracted_specific_location == 'unknown':
+                                #   extracted_specific_location = pop_code_to_specific_loc.get(pop_code_for_range, 'unknown')
+                                population_code_from_sl = pop_code_for_range
+                                method_used = "structured_lookup_heuristic_range_match"
+                                print(f"'{query_word}' not direct. Heuristic: Falls within range {query_prefix_upper}{start_num}-{query_prefix_upper}{end_num}.")
+                                break
+                            else:
+                                print(f"'{query_word}' heuristic match found, but country unknown. Will fall to RAG below.")
+
+    # Attempt 3: If primary query_word failed all structured lookups, try alternative_query_word (cleaned)
+    if method_used == 'unknown' and alternative_query_word_cleaned and alternative_query_word_cleaned != query_word:
+        print(f"'{query_word}' not found in structured (or heuristic). Trying alternative '{alternative_query_word_cleaned}'.")
+
+        # Try direct lookup for alternative word
+        structured_info_alt = final_structured_entries.get(alternative_query_word_cleaned.upper())
+        if structured_info_alt:
+          if extracted_country == 'unknown':
+            extracted_country = structured_info_alt['country']
+          if extracted_type == 'unknown':
+            extracted_type = structured_info_alt['type']
+          # if extracted_ethnicity == 'unknown':
+          #   extracted_ethnicity = structured_info_alt.get('ethnicity', 'unknown')
+          # if extracted_specific_location == 'unknown':
+          #   extracted_specific_location = structured_info_alt.get('specific_location', 'unknown')
+          population_code_from_sl = structured_info_alt['population_code']
+          method_used = "structured_lookup_alt_direct"
+          print(f"Alternative '{alternative_query_word_cleaned}' found in structured lookup (direct match).")
+        else:
+            # Try heuristic lookup for alternative word
+            alt_prefix, alt_num_str = _parse_individual_code_parts(alternative_query_word_cleaned)
+            if alt_prefix is not None and alt_num_str is not None:
+                try: alt_num = int(alt_num_str)
+                except ValueError: alt_num = None
+                if alt_num is not None:
+                    alt_prefix_upper = alt_prefix.upper()
+                    contiguous_ranges = master_structured_lookup.get('contiguous_ranges', defaultdict(list))
+                    pop_code_to_country = master_structured_lookup.get('pop_code_to_country', {})
+                    pop_code_to_ethnicity = master_structured_lookup.get('pop_code_to_ethnicity', {})
+                    pop_code_to_specific_loc = master_structured_lookup.get('pop_code_to_specific_loc', {})
+                    if alt_prefix_upper in contiguous_ranges:
+                        for start_num, end_num, pop_code_for_range in contiguous_ranges[alt_prefix_upper]:
+                            if start_num <= alt_num <= end_num:
+                                country_from_heuristic_alt = pop_code_to_country.get(pop_code_for_range, 'unknown')
+                                if country_from_heuristic_alt != 'unknown':
+                                  if extracted_country == 'unknown':
+                                    extracted_country = country_from_heuristic_alt
+                                  if extracted_type == 'unknown':
+                                    extracted_type = 'modern'
+                                  # if extracted_ethnicity == 'unknown':
+                                  #   extracted_ethnicity = pop_code_to_ethnicity.get(pop_code_for_range, 'unknown')
+                                  # if extracted_specific_location == 'unknown':
+                                  #   extracted_specific_location = pop_code_to_specific_loc.get(pop_code_for_range, 'unknown')
+                                  population_code_from_sl = pop_code_for_range
+                                  method_used = "structured_lookup_alt_heuristic_range_match"
+                                  break
+                                else:
+                                    print(f"Alternative '{alternative_query_word_cleaned}' heuristic match found, but country unknown. Will fall to RAG below.")
+
+    # use the context_for_llm to detect present_ancient before using llm model
+    # retrieved_chunks_text = []
+    # if document_chunks:
+    #   for idx in range(len(document_chunks)):
+    #           retrieved_chunks_text.append(document_chunks[idx])
+    # context_for_llm = ""
+    # all_context = "\n".join(retrieved_chunks_text) # 
+    # listOfcontexts = {"chunk": chunk, 
+    #         "all_output": all_output,
+    #         "document_chunk": all_context}
+    # label, context_for_llm = chooseContextLLM(listOfcontexts, query_word) 
+    # if not context_for_llm:
+    #   label, context_for_llm = chooseContextLLM(listOfcontexts, alternative_query_word_cleaned)
+    #   if not context_for_llm:
+    #     context_for_llm = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + extracted_features
+    # if context_for_llm:
+    #   extracted_type, explain = mtdna_classifier.detect_ancient_flag(context_for_llm)
+    #   extracted_type = extracted_type.lower()
+    #   sample_type_explanation = explain
+    # 5. Execute RAG if needed (either full RAG or targeted RAG for missing fields)
+
+    # Determine if a RAG call is necessary
+    run_rag = (extracted_country == 'unknown' or extracted_type == 'unknown')# or \
+               #extracted_ethnicity == 'unknown' or extracted_specific_location == 'unknown')
+    global_llm_model_for_counting_tokens = genai.GenerativeModel('gemini-1.5-flash-latest')
+    if run_rag:
+
+        # Determine the phrase for LLM query
+        rag_query_phrase = f"'{query_word}'"
+        if alternative_query_word_cleaned and alternative_query_word_cleaned != query_word:
+            rag_query_phrase += f" or its alternative word '{alternative_query_word_cleaned}'"
+
+        # Construct a more specific semantic query phrase for embedding if structured info is available
+        semantic_query_for_embedding = rag_query_phrase # Default
+        # if extracted_country != 'unknown': # If country is known from structured lookup (for targeted RAG)
+        #     if population_code_from_sl != 'unknown':
+        #         semantic_query_for_embedding = f"ethnicity and specific location for {query_word} population {population_code_from_sl} in {extracted_country}"
+        #     else: # If pop_code not found in structured, still use country hint
+        #         semantic_query_for_embedding = f"ethnicity and specific location for {query_word} in {extracted_country}"
+        # print(f"  DEBUG: Semantic query for embedding: '{semantic_query_for_embedding}'")
+
+
+        # Determine fields to ask LLM for and output format based on what's known/needed
+        prompt_instruction_prefix = ""
+        output_format_str = ""
+
+        # Determine if it's a full RAG or targeted RAG scenario based on what's already extracted
+        is_full_rag_scenario = True#(extracted_country == 'unknown')
+
+        if is_full_rag_scenario: # Full RAG scenario
+            output_format_str = "country_name, modern/ancient/unknown"#, ethnicity, specific_location/unknown"
+            method_used = "rag_llm"
+            print(f"Proceeding to FULL RAG for {rag_query_phrase}.")
+        # else: # Targeted RAG scenario (country/type already known, need ethnicity/specific_location)
+        #     if extracted_type == "unknown":
+        #         prompt_instruction_prefix = (
+        #             f"I already know the country is {extracted_country}. "
+        #             f"{f'The population code is {population_code_from_sl}. ' if population_code_from_sl != 'unknown' else ''}"
+        #         )
+        #         #output_format_str = "modern/ancient/unknown, ethnicity, specific_location/unknown"
+        #         output_format_str = "modern/ancient/unknown"
+        #     # else:
+        #     #     prompt_instruction_prefix = (
+        #     #         f"I already know the country is {extracted_country} and the sample type is {extracted_type}. "
+        #     #         f"{f'The population code is {population_code_from_sl}. ' if population_code_from_sl != 'unknown' else ''}"
+        #     #     )
+        #     #     output_format_str = "ethnicity, specific_location/unknown"
+
+        #     method_used = "hybrid_sl_rag"
+        #     print(f"Proceeding to TARGETED RAG for {rag_query_phrase}.")
+
+
+        # Calculate embedding cost for the primary query word
+        current_embedding_cost = 0
+        try:
+            query_embedding_vector = get_embedding(semantic_query_for_embedding, task_type="RETRIEVAL_QUERY")
+            query_embedding_tokens = global_llm_model_for_counting_tokens.count_tokens(semantic_query_for_embedding).total_tokens
+            current_embedding_cost += (query_embedding_tokens / 1000) * PRICE_PER_1K_EMBEDDING_INPUT
+            print(f"  DEBUG: Query embedding tokens (for '{semantic_query_for_embedding}'): {query_embedding_tokens}, cost: ${current_embedding_cost:.6f}")
+
+            if alternative_query_word_cleaned and alternative_query_word_cleaned != query_word:
+                alt_embedding_vector = get_embedding(alternative_query_word_cleaned, task_type="RETRIEVAL_QUERY")
+                alt_embedding_tokens = global_llm_model_for_counting_tokens.count_tokens(alternative_query_word_cleaned).total_tokens
+                current_embedding_cost += (alt_embedding_tokens / 1000) * PRICE_PER_1K_EMBEDDING_INPUT
+                print(f"  DEBUG: Alternative query ('{alternative_query_word_cleaned}') embedding tokens: {alt_embedding_tokens}, cost: ${current_embedding_cost:.6f}")
+
+        except Exception as e:
+            print(f"Error getting query embedding for RAG: {e}")
+            return extracted_country, extracted_type, "embedding_failed", extracted_ethnicity, extracted_specific_location, total_query_cost
+
+        if query_embedding_vector is None or query_embedding_vector.shape[0] == 0:
+            return extracted_country, extracted_type, "embedding_failed", extracted_ethnicity, extracted_specific_location, total_query_cost
+
+        D, I = faiss_index.search(np.array([query_embedding_vector]), 4)
+
+        retrieved_chunks_text = []
+        for idx in I[0]:
+            if 0 <= idx < len(document_chunks):
+                retrieved_chunks_text.append(document_chunks[idx])
+
+        context_for_llm = ""
+        
+        all_context = "\n".join(retrieved_chunks_text) # 
+        listOfcontexts = {"chunk": chunk, 
+            "all_output": all_output,
+            "document_chunk": all_context}
+        label, context_for_llm = chooseContextLLM(listOfcontexts, query_word) 
+        if not context_for_llm:
+          label, context_for_llm = chooseContextLLM(listOfcontexts, alternative_query_word_cleaned)
+          if not context_for_llm:
+            context_for_llm = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + extracted_features
+        #print("context for llm: ", label)    
+        # prompt_for_llm = (
+        #     f"{prompt_instruction_prefix}"
+        #     f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} or the mitochondrial DNA sample in general if these specific identifiers are not explicitly found. "
+        #     f"Identify its primary associated country/geographic location. "
+        #     f"Also, determine if the genetic sample or individual mentioned is from a 'modern' (present-day living individual) "
+        #     f"or 'ancient' (e.g., prehistoric remains, archaeological sample) source. "
+        #     f"If the text does not mention whether the sample is ancient or modern, assume the sample is modern unless otherwise explicitly described as ancient or archaeological. "
+        #     f"Additionally, extract its ethnicity and a more specific location (city/district level) within the predicted country. "
+        #     f"If any information is not explicitly present in the provided text snippets, state 'unknown' for that specific piece of information. "
+        #     f"Provide only the country, sample type, ethnicity, and specific location, do not add extra explanations.\n\n"
+        #     f"Text Snippets:\n{context_for_llm}\n\n"
+        #     f"Output Format: {output_format_str}"
+        # )
+        if len(context_for_llm) > 1000*1000:
+          context_for_llm = context_for_llm[:900000]
+        prompt_for_llm = (
+    f"{prompt_instruction_prefix}"
+    f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} or the mitochondrial DNA sample in general if these specific identifiers are not explicitly found. "
+    f"Identify its primary associated country/geographic location. "
+    f"Also, determine if the genetic sample or individual mentioned is from a 'modern' (present-day living individual) "
+    f"or 'ancient' (e.g., prehistoric remains, archaeological sample) source. "
+    f"If the text does not mention whether the sample is ancient or modern, assume the sample is modern unless otherwise explicitly described as ancient or archaeological. "
+    f"Provide only {output_format_str}. "
+    f"If any information is not explicitly present in the provided text snippets, state 'unknown' for that specific piece of information. "
+    f"If the country or sample type (modern/ancient) is not 'unknown', write 1 sentence after the output explaining how you inferred it from the text (one sentence for each)."
+    f"\n\nText Snippets:\n{context_for_llm}\n\n"
+    f"Output Format: {output_format_str}"
+)
+
+        llm_response_text, model_instance = call_llm_api(prompt_for_llm)
+        print("\n--- DEBUG INFO FOR RAG ---")
+        print("Retrieved Context Sent to LLM (first 500 chars):")
+        print(context_for_llm[:500] + "..." if len(context_for_llm) > 500 else context_for_llm)
+        print("\nRaw LLM Response:")
+        print(llm_response_text)
+        print("--- END DEBUG INFO ---")
+
+        llm_cost = 0
+        if model_instance:
+            try:
+                input_llm_tokens = global_llm_model_for_counting_tokens.count_tokens(prompt_for_llm).total_tokens
+                output_llm_tokens = global_llm_model_for_counting_tokens.count_tokens(llm_response_text).total_tokens
+                print(f"  DEBUG: LLM Input tokens: {input_llm_tokens}")
+                print(f"  DEBUG: LLM Output tokens: {output_llm_tokens}")
+                llm_cost = (input_llm_tokens / 1000) * PRICE_PER_1K_INPUT_LLM + \
+                           (output_llm_tokens / 1000) * PRICE_PER_1K_OUTPUT_LLM
+                print(f"  DEBUG: Estimated LLM cost: ${llm_cost:.6f}")
+            except Exception as e:
+                print(f"  DEBUG: Error counting LLM tokens: {e}")
+                llm_cost = 0
+
+        total_query_cost += current_embedding_cost + llm_cost
+        print(f"  DEBUG: Total estimated cost for this RAG query: ${total_query_cost:.6f}")
+        # Parse the LLM's response based on the Output Format actually used
+        # if output_format_str == "ethnicity, specific_location/unknown": # Targeted RAG output
+        #     extracted_ethnicity,extracted_specific_location = clean_llm_output(llm_response_text, output_format_str)
+        # elif output_format_str == "modern/ancient/unknown, ethnicity, specific_location/unknown":
+        #     extracted_type, extracted_ethnicity,extracted_specific_location=clean_llm_output(llm_response_text, output_format_str) 
+        # else: # Full RAG output (country, type, ethnicity, specific_location)
+        #     extracted_country,extracted_type, extracted_ethnicity,extracted_specific_location=clean_llm_output(llm_response_text, output_format_str) 
+        metadata_list = parse_multi_sample_llm_output(llm_response_text, output_format_str)
+        merge_metadata = merge_metadata_outputs(metadata_list)
+        if output_format_str == "country_name, modern/ancient/unknown":
+          extracted_country, extracted_type = merge_metadata["country"], merge_metadata["sample_type"]  
+          country_explanation,sample_type_explanation = merge_metadata["country_explanation"], merge_metadata["sample_type_explanation"]
+        elif output_format_str == "modern/ancient/unknown":
+          extracted_type = merge_metadata["sample_type"]  
+          sample_type_explanation = merge_metadata["sample_type_explanation"]  
+    # 6. Optional: Second LLM call for specific_location from general knowledge if still unknown
+    # if extracted_specific_location == 'unknown':
+    #     # Check if we have enough info to ask general knowledge LLM
+    #     if extracted_country != 'unknown' and extracted_ethnicity != 'unknown':
+    #         print(f"  DEBUG: Specific location still unknown. Querying general knowledge LLM from '{extracted_ethnicity}' and '{extracted_country}'.")
+
+    #         general_knowledge_prompt = (
+    #             f"Based on general knowledge, what is a highly specific location (city or district) "
+    #             f"associated with the ethnicity '{extracted_ethnicity}' in '{extracted_country}'? "
+    #             f"Consider the context of scientific studies on human genetics, if known. "
+    #             f"If no common specific location is known, state 'unknown'. "
+    #             f"Provide only the city or district name, or 'unknown'."
+    #         )
+
+    #         general_llm_response, general_llm_model_instance = call_llm_api(general_knowledge_prompt, model_name='gemini-1.5-flash-latest')
+
+    #         if general_llm_response and general_llm_response.lower().strip() != 'unknown':
+    #             extracted_specific_location = general_llm_response.strip() + " (predicted from general knowledge)"
+    #             # Add cost of this second LLM call
+    #             if general_llm_model_instance:
+    #                 try:
+    #                     gk_input_tokens = general_llm_model_instance.count_tokens(general_knowledge_prompt).total_tokens
+    #                     gk_output_tokens = general_llm_model_instance.count_tokens(general_llm_response).total_tokens
+    #                     gk_cost = (gk_input_tokens / 1000) * PRICE_PER_1K_INPUT_LLM + \
+    #                               (gk_output_tokens / 1000) * PRICE_PER_1K_OUTPUT_LLM
+    #                     print(f"  DEBUG: General Knowledge LLM cost to predict specific location alone: ${gk_cost:.6f}")
+    #                     total_query_cost += gk_cost # Accumulate cost
+    #                 except Exception as e:
+    #                     print(f"  DEBUG: Error counting GK LLM tokens: {e}")
+    #         else:
+    #             print("  DEBUG: General knowledge LLM returned unknown or empty for specific location.")
+    # # 6. Optional: Second LLM call for ethnicity from general knowledge if still unknown
+    # if extracted_ethnicity == 'unknown':
+    #     # Check if we have enough info to ask general knowledge LLM
+    #     if extracted_country != 'unknown' and extracted_specific_location != 'unknown':
+    #         print(f"  DEBUG: Ethnicity still unknown. Querying general knowledge LLM from '{extracted_specific_location}' and '{extracted_country}'.")
+
+    #         general_knowledge_prompt = (
+    #             f"Based on general knowledge, what is a highly ethnicity (population) "
+    #             f"associated with the specific location '{extracted_specific_location}' in '{extracted_country}'? "
+    #             f"Consider the context of scientific studies on human genetics, if known. "
+    #             f"If no common ethnicity is known, state 'unknown'. "
+    #             f"Provide only the ethnicity or popluation name, or 'unknown'."
+    #         )
+
+    #         general_llm_response, general_llm_model_instance = call_llm_api(general_knowledge_prompt, model_name='gemini-1.5-flash-latest')
+
+    #         if general_llm_response and general_llm_response.lower().strip() != 'unknown':
+    #             extracted_ethnicity = general_llm_response.strip() + " (predicted from general knowledge)"
+    #             # Add cost of this second LLM call
+    #             if general_llm_model_instance:
+    #                 try:
+    #                     gk_input_tokens = general_llm_model_instance.count_tokens(general_knowledge_prompt).total_tokens
+    #                     gk_output_tokens = general_llm_model_instance.count_tokens(general_llm_response).total_tokens
+    #                     gk_cost = (gk_input_tokens / 1000) * PRICE_PER_1K_INPUT_LLM + \
+    #                               (gk_output_tokens / 1000) * PRICE_PER_1K_OUTPUT_LLM
+    #                     print(f"  DEBUG: General Knowledge LLM cost to predict ethnicity alone: ${gk_cost:.6f}")
+    #                     total_query_cost += gk_cost # Accumulate cost
+    #                 except Exception as e:
+    #                     print(f"  DEBUG: Error counting GK LLM tokens: {e}")
+    #         else:
+    #             print("  DEBUG: General knowledge LLM returned unknown or empty for ethnicity.")            
+
+
+    #return extracted_country, extracted_type, method_used, extracted_ethnicity, extracted_specific_location, total_query_cost
+    return extracted_country, extracted_type, method_used, country_explanation, sample_type_explanation, total_query_cost

mtdna_backend.py

@@ -3,7 +3,9 @@ from collections import Counter
 import csv
 import os
 from functools import lru_cache
+import mtdna_ui_app
 from mtdna_classifier import classify_sample_location 
+from iterate3 import data_preprocess, model, pipeline
 import subprocess
 import json
 import pandas as pd
@@ -13,41 +15,47 @@ import tempfile
 import gspread
 from oauth2client.service_account import ServiceAccountCredentials
 from io import StringIO
+import hashlib
+import threading
 
-@lru_cache(maxsize=128)
-def classify_sample_location_cached(accession):
-    return classify_sample_location(accession)
+# @lru_cache(maxsize=3600)
+# def classify_sample_location_cached(accession):
+#     return classify_sample_location(accession)
+
+@lru_cache(maxsize=3600)
+def pipeline_classify_sample_location_cached(accession):
+    return pipeline.pipeline_with_gemini([accession])    
 
 # Count and suggest final location
-def compute_final_suggested_location(rows):
-    candidates = [
-        row.get("Predicted Location", "").strip()
-        for row in rows
-        if row.get("Predicted Location", "").strip().lower() not in ["", "sample id not found", "unknown"]
-    ] + [
-        row.get("Inferred Region", "").strip()
-        for row in rows
-        if row.get("Inferred Region", "").strip().lower() not in  ["", "sample id not found", "unknown"]
-    ]
-
-    if not candidates:
-        return Counter(), ("Unknown", 0)
-    # Step 1: Combine into one string and split using regex to handle commas, line breaks, etc.
-    tokens = []
-    for item in candidates:
-        # Split by comma, whitespace, and newlines
-        parts = re.split(r'[\s,]+', item)
-        tokens.extend(parts)
-
-    # Step 2: Clean and normalize tokens
-    tokens = [word.strip() for word in tokens if word.strip().isalpha()]  # Keep only alphabetic tokens
-
-    # Step 3: Count
-    counts = Counter(tokens)
-
-    # Step 4: Get most common
-    top_location, count = counts.most_common(1)[0]
-    return counts, (top_location, count)
+# def compute_final_suggested_location(rows):
+#     candidates = [
+#         row.get("Predicted Location", "").strip()
+#         for row in rows
+#         if row.get("Predicted Location", "").strip().lower() not in ["", "sample id not found", "unknown"]
+#     ] + [
+#         row.get("Inferred Region", "").strip()
+#         for row in rows
+#         if row.get("Inferred Region", "").strip().lower() not in  ["", "sample id not found", "unknown"]
+#     ]
+
+#     if not candidates:
+#         return Counter(), ("Unknown", 0)
+#     # Step 1: Combine into one string and split using regex to handle commas, line breaks, etc.
+#     tokens = []
+#     for item in candidates:
+#         # Split by comma, whitespace, and newlines
+#         parts = re.split(r'[\s,]+', item)
+#         tokens.extend(parts)
+
+#     # Step 2: Clean and normalize tokens
+#     tokens = [word.strip() for word in tokens if word.strip().isalpha()]  # Keep only alphabetic tokens
+
+#     # Step 3: Count
+#     counts = Counter(tokens)
+
+#     # Step 4: Get most common
+#     top_location, count = counts.most_common(1)[0]
+#     return counts, (top_location, count)
 
 # Store feedback (with required fields)
 
@@ -100,74 +108,216 @@ def extract_accessions_from_input(file=None, raw_text=""):
                 seen.add(acc)
 
     return list(accessions), None
+# ✅ Add a new helper to backend: `filter_unprocessed_accessions()`
+def get_incomplete_accessions(file_path):
+    df = pd.read_excel(file_path)
+
+    incomplete_accessions = []
+    for _, row in df.iterrows():
+        sample_id = str(row.get("Sample ID", "")).strip()
 
-def summarize_results(accession):
+        # Skip if no sample ID
+        if not sample_id:
+            continue
+
+        # Drop the Sample ID and check if the rest is empty
+        other_cols = row.drop(labels=["Sample ID"], errors="ignore")
+        if other_cols.isna().all() or (other_cols.astype(str).str.strip() == "").all():
+            # Extract the accession number from the sample ID using regex
+            match = re.search(r"\b[A-Z]{2,4}\d{4,}", sample_id)
+            if match:
+                incomplete_accessions.append(match.group(0))
+    print(len(incomplete_accessions))
+    return incomplete_accessions
+
+def summarize_results(accession, KNOWN_OUTPUT_PATH = "/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/iterate3/known_samples.xlsx"):
+    # try cache first
+    cached = check_known_output(accession)
+    if cached:
+        print(f"✅ Using cached result for {accession}")
+        return [[
+            cached["Sample ID"],
+            cached["Predicted Country"],
+            cached["Country Explanation"],
+            cached["Predicted Sample Type"],
+            cached["Sample Type Explanation"],
+            cached["Sources"],
+            cached["Time cost"]
+        ]]
+    # only run when nothing in the cache  
     try:
-        output, labelAncient_Modern, explain_label = classify_sample_location_cached(accession)
-        #print(output)
+        outputs = pipeline_classify_sample_location_cached(accession)
+        # outputs = {'KU131308': {'isolate':'BRU18',
+        # 'country': {'brunei': ['ncbi',
+        # 'rag_llm-The text mentions "BRU18 Brunei Borneo" in a table listing various samples, and it is not described as ancient or archaeological.']},
+        # 'sample_type': {'modern': 
+        # ['rag_llm-The text mentions "BRU18 Brunei Borneo" in a table listing various samples, and it is not described as ancient or archaeological.']},
+        # 'query_cost': 9.754999999999999e-05,
+        # 'time_cost': '24.776 seconds',
+        # 'source': ['https://doi.org/10.1007/s00439-015-1620-z',
+        # 'https://static-content.springer.com/esm/art%3A10.1007%2Fs00439-015-1620-z/MediaObjects/439_2015_1620_MOESM1_ESM.pdf',
+        # 'https://static-content.springer.com/esm/art%3A10.1007%2Fs00439-015-1620-z/MediaObjects/439_2015_1620_MOESM2_ESM.xls']}}
     except Exception as e:
-        return [], f"Error: {e}", f"Error: {e}", f"Error: {e}"
+        return []#, f"Error: {e}", f"Error: {e}", f"Error: {e}"
 
-    if accession not in output:
-        return [], "Accession not found in results.", "Accession not found in results.", "Accession not found in results."
+    if accession not in outputs:
+        return []#, "Accession not found in results.", "Accession not found in results.", "Accession not found in results."
 
-    isolate = next((k for k in output if k != accession), None)
     row_score = []
     rows = []
+    save_rows = []
+    for key in outputs:
+        pred_country, pred_sample, country_explanation, sample_explanation = "unknown","unknown","unknown","unknown" 
+        for section, results in outputs[key].items():
+          if section == "country" or section =="sample_type":
+            pred_output = "\n".join(list(results.keys()))
+            output_explanation = ""
+            for result, content in results.items():
+              if len(result) == 0:  result = "unknown"
+              if len(content) == 0: output_explanation = "unknown"
+              else:
+                output_explanation += 'Method: ' + "\nMethod: ".join(content)  + "\n"
+            if section == "country":
+              pred_country, country_explanation = pred_output, output_explanation
+            elif section == "sample_type":
+              pred_sample, sample_explanation = pred_output, output_explanation    
+          if outputs[key]["isolate"].lower()!="unknown":
+            label = key + "(Isolate: " + outputs[key]["isolate"] + ")"
+          else: label = key  
+        if len(outputs[key]["source"]) == 0:  outputs[key]["source"] = ["No Links"]
+        row = {
+            "Sample ID": label,
+            "Predicted Country": pred_country,
+            "Country Explanation": country_explanation,
+            "Predicted Sample Type":pred_sample,
+            "Sample Type Explanation":sample_explanation,
+            "Sources": "\n".join(outputs[key]["source"]),
+            "Time cost": outputs[key]["time_cost"]
+        }
+        #row_score.append(row)
+        rows.append(list(row.values()))
+
+        save_row = {
+            "Sample ID": label,
+            "Predicted Country": pred_country,
+            "Country Explanation": country_explanation,
+            "Predicted Sample Type":pred_sample,
+            "Sample Type Explanation":sample_explanation,
+            "Sources": "\n".join(outputs[key]["source"]),
+            "Query_cost": outputs[key]["query_cost"],
+            "Time cost": outputs[key]["time_cost"]
+        }
+        #row_score.append(row)
+        save_rows.append(list(save_row.values()))
+
+    # #location_counts, (final_location, count) = compute_final_suggested_location(row_score)
+    # summary_lines = [f"### 🧭 Location Summary:\n"]
+    # summary_lines += [f"- **{loc}**: {cnt} times" for loc, cnt in location_counts.items()]
+    # summary_lines.append(f"\n**Final Suggested Location:** 🗺️ **{final_location}** (mentioned {count} times)")
+    # summary = "\n".join(summary_lines)
+
+    # save the new running sample to known excel file
+    try:
+      df_new = pd.DataFrame(save_rows, columns=["Sample ID", "Predicted Country", "Country Explanation", "Predicted Sample Type", "Sample Type Explanation", "Sources", "Query_cost","Time cost"])
+      if os.path.exists(KNOWN_OUTPUT_PATH):
+          df_old = pd.read_excel(KNOWN_OUTPUT_PATH)
+          df_combined = pd.concat([df_old, df_new]).drop_duplicates(subset="Sample ID")
+      else:
+          df_combined = df_new
+      df_combined.to_excel(KNOWN_OUTPUT_PATH, index=False)
+    except Exception as e:
+      print(f"⚠️ Failed to save known output: {e}")
 
-    for key in [accession, isolate]:
-        if key not in output:
-            continue
-        sample_id_label = f"{key} ({'accession number' if key == accession else 'isolate of accession'})"
-        for section, techniques in output[key].items():
-            for technique, content in techniques.items():
-                source = content.get("source", "")
-                predicted = content.get("predicted_location", "")
-                haplogroup = content.get("haplogroup", "")
-                inferred = content.get("inferred_location", "")
-                context = content.get("context_snippet", "")[:300] if "context_snippet" in content else ""
-    
-                row = {
-                    "Sample ID": sample_id_label,
-                    "Technique": technique,
-                    "Source": f"The region of haplogroup is inferred\nby using this source: {source}" if technique == "haplogroup" else source,
-                    "Predicted Location": "" if technique == "haplogroup" else predicted,
-                    "Haplogroup": haplogroup if technique == "haplogroup" else "",
-                    "Inferred Region": inferred if technique == "haplogroup" else "",
-                    "Context Snippet": context
-                }
-
-                row_score.append(row)
-                rows.append(list(row.values()))
-
-    location_counts, (final_location, count) = compute_final_suggested_location(row_score)
-    summary_lines = [f"### 🧭 Location Frequency Summary", "After counting all predicted and inferred locations:\n"]
-    summary_lines += [f"- **{loc}**: {cnt} times" for loc, cnt in location_counts.items()]
-    summary_lines.append(f"\n**Final Suggested Location:** 🗺️ **{final_location}** (mentioned {count} times)")
-    summary = "\n".join(summary_lines)
-    return rows, summary, labelAncient_Modern, explain_label
+    return rows#, summary, labelAncient_Modern, explain_label
 
 # save the batch input in excel file
-def save_to_excel(all_rows, summary_text, flag_text, filename):
-    with pd.ExcelWriter(filename) as writer:
-        # Save table
-        df = pd.DataFrame(all_rows, columns=["Sample ID", "Technique", "Source", "Predicted Location", "Haplogroup", "Inferred Region", "Context Snippet"])
-        df.to_excel(writer, sheet_name="Detailed Results", index=False)
-        
-        # Save summary
-        summary_df = pd.DataFrame({"Summary": [summary_text]})
-        summary_df.to_excel(writer, sheet_name="Summary", index=False)
+# def save_to_excel(all_rows, summary_text, flag_text, filename):
+#     with pd.ExcelWriter(filename) as writer:
+#         # Save table
+#         df_new = pd.DataFrame(all_rows, columns=["Sample ID", "Predicted Country", "Country Explanation", "Predicted Sample Type", "Sample Type Explanation", "Sources", "Time cost"])
+#         df.to_excel(writer, sheet_name="Detailed Results", index=False)
+#         try:
+#           df_old = pd.read_excel(filename)
+#         except:
+#           df_old = pd.DataFrame([[]], columns=["Sample ID", "Predicted Country", "Country Explanation", "Predicted Sample Type", "Sample Type Explanation", "Sources", "Time cost"])
+#         df_combined = pd.concat([df_old, df_new]).drop_duplicates(subset="Sample ID")
+#         # if os.path.exists(filename):
+#         #   df_old = pd.read_excel(filename)
+#         #   df_combined = pd.concat([df_old, df_new]).drop_duplicates(subset="Sample ID")
+#         # else:
+#         #     df_combined = df_new
+#         df_combined.to_excel(filename, index=False)
+#           # # Save summary
+#         # summary_df = pd.DataFrame({"Summary": [summary_text]})
+#         # summary_df.to_excel(writer, sheet_name="Summary", index=False)
         
-        # Save flag
-        flag_df = pd.DataFrame({"Flag": [flag_text]})
-        flag_df.to_excel(writer, sheet_name="Ancient_Modern_Flag", index=False)
+#         # # Save flag
+#         # flag_df = pd.DataFrame({"Flag": [flag_text]})
+#         # flag_df.to_excel(writer, sheet_name="Ancient_Modern_Flag", index=False)
+# def save_to_excel(all_rows, summary_text, flag_text, filename):
+#     df_new = pd.DataFrame(all_rows, columns=[
+#         "Sample ID", "Predicted Country", "Country Explanation", 
+#         "Predicted Sample Type", "Sample Type Explanation", 
+#         "Sources", "Time cost"
+#     ])
+
+#     try:
+#         if os.path.exists(filename):
+#             df_old = pd.read_excel(filename)
+#         else:
+#             df_old = pd.DataFrame(columns=df_new.columns)
+#     except Exception as e:
+#         print(f"⚠️ Warning reading old Excel file: {e}")
+#         df_old = pd.DataFrame(columns=df_new.columns)
+
+#     #df_combined = pd.concat([df_new, df_old], ignore_index=True).drop_duplicates(subset="Sample ID", keep="first")
+#     df_old.set_index("Sample ID", inplace=True)
+#     df_new.set_index("Sample ID", inplace=True)
+
+#     df_old.update(df_new)  # <-- update matching rows in df_old with df_new content
+
+#     df_combined = df_old.reset_index()
+
+#     try:
+#         df_combined.to_excel(filename, index=False)
+#     except Exception as e:
+#         print(f"❌ Failed to write Excel file {filename}: {e}")
+def save_to_excel(all_rows, summary_text, flag_text, filename, is_resume=False):
+    df_new = pd.DataFrame(all_rows, columns=[
+        "Sample ID", "Predicted Country", "Country Explanation", 
+        "Predicted Sample Type", "Sample Type Explanation", 
+        "Sources", "Time cost"
+    ])
+
+    if is_resume and os.path.exists(filename):
+        try:
+            df_old = pd.read_excel(filename)
+        except Exception as e:
+            print(f"⚠️ Warning reading old Excel file: {e}")
+            df_old = pd.DataFrame(columns=df_new.columns)
+
+        # Set index and update existing rows
+        df_old.set_index("Sample ID", inplace=True)
+        df_new.set_index("Sample ID", inplace=True)
+        df_old.update(df_new)
+
+        df_combined = df_old.reset_index()
+    else:
+        # If not resuming or file doesn't exist, just use new rows
+        df_combined = df_new
+
+    try:
+        df_combined.to_excel(filename, index=False)
+    except Exception as e:
+        print(f"❌ Failed to write Excel file {filename}: {e}")
+
 
 # save the batch input in JSON file
 def save_to_json(all_rows, summary_text, flag_text, filename):
     output_dict = {
-        "Detailed_Results": all_rows,  # <-- make sure this is a plain list, not a DataFrame
-        "Summary_Text": summary_text,
-        "Ancient_Modern_Flag": flag_text
+        "Detailed_Results": all_rows#,  # <-- make sure this is a plain list, not a DataFrame
+        # "Summary_Text": summary_text,
+        # "Ancient_Modern_Flag": flag_text
     }
 
     # If all_rows is a DataFrame, convert it
@@ -189,13 +339,13 @@ def save_to_txt(all_rows, summary_text, flag_text, filename):
         f.write("=== Detailed Results ===\n")
         f.write(output + "\n")
 
-        f.write("\n=== Summary ===\n")
-        f.write(summary_text + "\n")
+        # f.write("\n=== Summary ===\n")
+        # f.write(summary_text + "\n")
         
-        f.write("\n=== Ancient/Modern Flag ===\n")
-        f.write(flag_text + "\n")
+        # f.write("\n=== Ancient/Modern Flag ===\n")
+        # f.write(flag_text + "\n")
 
-def save_batch_output(all_rows, summary_text, flag_text, output_type):
+def save_batch_output(all_rows, output_type, summary_text=None, flag_text=None):
     tmp_dir = tempfile.mkdtemp()
 
     #html_table = all_rows.value  # assuming this is stored somewhere
@@ -219,34 +369,142 @@ def save_batch_output(all_rows, summary_text, flag_text, output_type):
         return gr.update(visible=False)  # invalid option
     
     return gr.update(value=file_path, visible=True)
+# save cost by checking the known outputs
+
+def check_known_output(accession, KNOWN_OUTPUT_PATH = "/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/iterate3/known_samples.xlsx"):
+    if not os.path.exists(KNOWN_OUTPUT_PATH):
+        return None
+
+    try:
+        df = pd.read_excel(KNOWN_OUTPUT_PATH)
+        match = re.search(r"\b[A-Z]{2,4}\d{4,}", accession)
+        if match:
+          accession = match.group(0)
+          
+        matched = df[df["Sample ID"].str.contains(accession, case=False, na=False)]
+        if not matched.empty:
+            return matched.iloc[0].to_dict()  # Return the cached row
+    except Exception as e:
+        print(f"⚠️ Failed to load known samples: {e}")
+        return None
+
+USER_USAGE_TRACK_FILE = "/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/iterate3/user_usage_log.json"
+
+def hash_user_id(user_input):
+    return hashlib.sha256(user_input.encode()).hexdigest()
+
+# ✅ Load and save usage count
+
+# def load_user_usage():
+#     if os.path.exists(USER_USAGE_TRACK_FILE):
+#         with open(USER_USAGE_TRACK_FILE, "r") as f:
+#             return json.load(f)
+#     return {}
+
+def load_user_usage():
+    if not os.path.exists(USER_USAGE_TRACK_FILE):
+        return {}
+
+    try:
+        with open(USER_USAGE_TRACK_FILE, "r") as f:
+            content = f.read().strip()
+            if not content:
+                return {}  # file is empty
+            return json.loads(content)
+    except (json.JSONDecodeError, ValueError):
+        print("⚠️ Warning: user_usage.json is corrupted or invalid. Resetting.")
+        return {}  # fallback to empty dict
+
+
+def save_user_usage(usage):
+    with open(USER_USAGE_TRACK_FILE, "w") as f:
+        json.dump(usage, f, indent=2)
+
+# def increment_usage(user_id, num_samples=1):
+#     usage = load_user_usage()
+#     if user_id not in usage:
+#         usage[user_id] = 0
+#     usage[user_id] += num_samples
+#     save_user_usage(usage)
+#     return usage[user_id]
+def increment_usage(email: str, count: int):
+    usage = load_user_usage()
+    email_key = email.strip().lower()
+    usage[email_key] = usage.get(email_key, 0) + count
+    save_user_usage(usage)
+    return usage[email_key]
 
 # run the batch
-def summarize_batch(file=None, raw_text=""):
+def summarize_batch(file=None, raw_text="", resume_file=None, user_email="", 
+                    stop_flag=None, output_file_path=None, 
+                    limited_acc=50, yield_callback=None):
+    if user_email:
+      limited_acc += 10
     accessions, error = extract_accessions_from_input(file, raw_text)
     if error:
-        return [], "", "", f"Error: {error}"
+        #return [], "", "", f"Error: {error}"
+        return [], f"Error: {error}", 0, "", ""
+    if resume_file:
+      accessions = get_incomplete_accessions(resume_file)
+    tmp_dir = tempfile.mkdtemp()
+    if not output_file_path:
+      if resume_file:
+        output_file_path = os.path.join(tmp_dir, resume_file) 
+      else:  
+        output_file_path = os.path.join(tmp_dir, "batch_output_live.xlsx")
 
     all_rows = []
-    all_summaries = []
-    all_flags = []
-
-    for acc in accessions:
+    # all_summaries = []
+    # all_flags = []
+    progress_lines = []
+    warning = ""
+    if len(accessions) > limited_acc:  
+      accessions = accessions[:limited_acc]
+      warning = f"Your number of accessions is more than the {limited_acc}, only handle first {limited_acc} accessions"
+    for i, acc in enumerate(accessions):
+        if stop_flag and stop_flag.value:
+            line = f"🛑 Stopped at {acc} ({i+1}/{len(accessions)})"
+            progress_lines.append(line)
+            if yield_callback:
+              yield_callback(line)
+            print("🛑 User requested stop.")
+            break
+        print(f"[{i+1}/{len(accessions)}] Processing {acc}")
         try:
-            rows, summary, label, explain = summarize_results(acc)
+            # rows, summary, label, explain = summarize_results(acc)
+            rows = summarize_results(acc)
             all_rows.extend(rows)
-            all_summaries.append(f"**{acc}**\n{summary}")
-            all_flags.append(f"**{acc}**\n### 🏺 Ancient/Modern Flag\n**{label}**\n\n_Explanation:_ {explain}")
+            # all_summaries.append(f"**{acc}**\n{summary}")
+            # all_flags.append(f"**{acc}**\n### 🏺 Ancient/Modern Flag\n**{label}**\n\n_Explanation:_ {explain}")
+            #save_to_excel(all_rows, summary_text="", flag_text="", filename=output_file_path)
+            save_to_excel(all_rows, summary_text="", flag_text="", filename=output_file_path, is_resume=bool(resume_file))
+            line = f"✅ Processed {acc} ({i+1}/{len(accessions)})"
+            progress_lines.append(line)
+            if yield_callback:
+              yield_callback(f"✅ Processed {acc} ({i+1}/{len(accessions)})")
         except Exception as e:
-            all_summaries.append(f"**{acc}**: Failed - {e}")
-
+            print(f"❌ Failed to process {acc}: {e}")
+            continue
+            #all_summaries.append(f"**{acc}**: Failed - {e}")
+        #progress_lines.append(f"✅ Processed {acc} ({i+1}/{len(accessions)})")
+        limited_acc -= 1
     """for row in all_rows:
           source_column = row[2]  # Assuming the "Source" is in the 3rd column (index 2)
           
           if source_column.startswith("http"):  # Check if the source is a URL
               # Wrap it with HTML anchor tags to make it clickable
               row[2] = f'<a href="{source_column}" target="_blank" style="color: blue; text-decoration: underline;">{source_column}</a>'"""
-      
-
-    summary_text = "\n\n---\n\n".join(all_summaries)
-    flag_text = "\n\n---\n\n".join(all_flags)
-    return all_rows, summary_text, flag_text, gr.update(visible=True), gr.update(visible=False)
+    if not warning:
+      warning = f"You only have {limited_acc} left"
+    if user_email.strip():
+        user_hash = hash_user_id(user_email)
+        total_queries = increment_usage(user_hash, len(all_rows))
+    else:
+        total_queries = 0
+    yield_callback("✅ Finished!")
+
+    # summary_text = "\n\n---\n\n".join(all_summaries)
+    # flag_text = "\n\n---\n\n".join(all_flags)
+    #return all_rows, summary_text, flag_text, gr.update(visible=True), gr.update(visible=False)
+    #return all_rows, gr.update(visible=True), gr.update(visible=False)
+    return all_rows, output_file_path, total_queries, "\n".join(progress_lines), warning

mtdna_classifier.py

@@ -1,524 +1,707 @@
-# mtDNA Location Classifier MVP (Google Colab)
-# Accepts accession number → Fetches PubMed ID + isolate name → Gets abstract → Predicts location
-import os
-import subprocess
-import re
-from Bio import Entrez
-import fitz
-import spacy
-from spacy.cli import download
-from NER.PDF import pdf
-from NER.WordDoc import wordDoc
-from NER.html import extractHTML
-from NER.word2Vec import word2vec
-from transformers import pipeline
-import urllib.parse, requests
-from pathlib import Path
-from upgradeClassify import filter_context_for_sample, infer_location_for_sample
-# Set your email (required by NCBI Entrez)
-#Entrez.email = "[email protected]"
-import nltk
-
-nltk.download("stopwords")
-#nltk.download("punkt")
-nltk.download('punkt', download_dir='/home/user/nltk_data')
-
-nltk.download('punkt_tab')
-# Step 1: Get PubMed ID from Accession using EDirect
-
-'''def get_info_from_accession(accession):
-    cmd = f'{os.environ["HOME"]}/edirect/esummary -db nuccore -id {accession} -format medline | egrep "PUBMED|isolate"'
-    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-    output = result.stdout
-    pubmedID, isolate = "", ""
-    for line in output.split("\n"):
-      if len(line) > 0:
-        if "PUBMED" in line:
-          pubmedID = line.split()[-1]
-        if "isolate" in line:  # Check for isolate information
-          # Try direct GenBank annotation: /isolate="XXX"
-          match1 = re.search(r'/isolate\s*=\s*"([^"]+)"', line)  # search on current line
-          if match1:
-            isolate = match1.group(1)
-          else:
-            # Try from DEFINITION line: ...isolate XXX...
-            match2 = re.search(r'isolate\s+([A-Za-z0-9_-]+)', line) # search on current line
-            if match2:
-              isolate = match2.group(1)'''
-from Bio import Entrez, Medline
-import re
-
-Entrez.email = "[email protected]"
-
-def get_info_from_accession(accession):
-    try:
-        handle = Entrez.efetch(db="nuccore", id=accession, rettype="medline", retmode="text")
-        text = handle.read()
-        handle.close()
-
-        # Extract PUBMED ID from the Medline text
-        pubmed_match = re.search(r'PUBMED\s+(\d+)', text)
-        pubmed_id = pubmed_match.group(1) if pubmed_match else ""
-
-        # Extract isolate if available
-        isolate_match = re.search(r'/isolate="([^"]+)"', text)
-        if not isolate_match:
-            isolate_match = re.search(r'isolate\s+([A-Za-z0-9_-]+)', text)
-        isolate = isolate_match.group(1) if isolate_match else ""
-
-        if not pubmed_id:
-            print(f"⚠️ No PubMed ID found for accession {accession}")
-
-        return pubmed_id, isolate
-
-    except Exception as e:
-        print("❌ Entrez error:", e)
-        return "", ""
-# Step 2: Get doi link to access the paper
-'''def get_doi_from_pubmed_id(pubmed_id):
-    cmd = f'{os.environ["HOME"]}/edirect/esummary -db pubmed -id {pubmed_id} -format medline | grep -i "AID"'
-    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-    output = result.stdout
-
-    doi_pattern = r'10\.\d{4,9}/[-._;()/:A-Z0-9]+(?=\s*\[doi\])'
-    match = re.search(doi_pattern, output, re.IGNORECASE)
-
-    if match:
-        return match.group(0)
-    else:
-        return None  # or raise an Exception with a helpful message'''
-
-def get_doi_from_pubmed_id(pubmed_id):
-    try:
-        handle = Entrez.efetch(db="pubmed", id=pubmed_id, rettype="medline", retmode="text")
-        records = list(Medline.parse(handle))
-        handle.close()
-
-        if not records:
-            return None
-        
-        record = records[0]
-        if "AID" in record:
-            for aid in record["AID"]:
-                if "[doi]" in aid:
-                    return aid.split(" ")[0]  # extract the DOI
-
-        return None
-
-    except Exception as e:
-        print(f"❌ Failed to get DOI from PubMed ID {pubmed_id}: {e}")
-        return None
-
-
-# Step 3: Extract Text: Get the paper (html text), sup. materials (pdf, doc, excel) and do text-preprocessing
-# Step 3.1: Extract Text
-# sub: download excel file
-def download_excel_file(url, save_path="temp.xlsx"):
-    if "view.officeapps.live.com" in url:
-        parsed_url = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
-        real_url = urllib.parse.unquote(parsed_url["src"][0])
-        response = requests.get(real_url)
-        with open(save_path, "wb") as f:
-            f.write(response.content)
-        return save_path
-    elif url.startswith("http") and (url.endswith(".xls") or url.endswith(".xlsx")):
-        response = requests.get(url)
-        response.raise_for_status()  # Raises error if download fails
-        with open(save_path, "wb") as f:
-            f.write(response.content)
-        return save_path
-    else:
-        print("URL must point directly to an .xls or .xlsx file\n or it already downloaded.")
-        return url
-def get_paper_text(doi,id,manualLinks=None):
-  # create the temporary folder to contain the texts
-  '''folder_path = Path("data/"+str(id))
-  if not folder_path.exists():
-      cmd = f'mkdir data/{id}'
-      result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-      print("data/"+str(id) +" created.")
-  else:
-      print("data/"+str(id) +" already exists.")'''
-
-  cmd = f'mkdir data/{id}'
-  result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)  
-  saveLinkFolder = "data/"+id
-
-  link = 'https://doi.org/' + doi
-  '''textsToExtract = { "doiLink":"paperText"
-                        "file1.pdf":"text1",
-                        "file2.doc":"text2",
-                        "file3.xlsx":excelText3'''
-  textsToExtract = {}
-  # get the file to create listOfFile for each id
-  html = extractHTML.HTML("",link)
-  jsonSM = html.getSupMaterial()
-  text = ""
-  links  = [link] + sum((jsonSM[key] for key in jsonSM),[])
-  if manualLinks != None:
-    links += manualLinks
-  for l in links:
-    # get the main paper
-    name = l.split("/")[-1]
-    #file_path = folder_path / name
-    if l == link:
-      text = html.getListSection()
-      textsToExtract[link] = text
-    elif l.endswith(".pdf"):
-      '''if file_path.is_file():
-          l = saveLinkFolder + "/" + name
-          print("File exists.")'''
-      p = pdf.PDF(l,saveLinkFolder,doi)
-      f = p.openPDFFile()
-      pdf_path = saveLinkFolder + "/" + l.split("/")[-1]
-      doc = fitz.open(pdf_path)
-      text = "\n".join([page.get_text() for page in doc])
-      textsToExtract[l] = text
-    elif l.endswith(".doc") or l.endswith(".docx"):
-      d = wordDoc.wordDoc(l,saveLinkFolder)
-      text = d.extractTextByPage()
-      textsToExtract[l] = text
-    elif l.split(".")[-1].lower() in "xlsx":
-      wc = word2vec.word2Vec()
-      # download excel file if it not downloaded yet
-      savePath = saveLinkFolder +"/"+ l.split("/")[-1]
-      excelPath = download_excel_file(l, savePath)
-      corpus = wc.tableTransformToCorpusText([],excelPath)
-      text = ''
-      for c in corpus:
-        para = corpus[c]
-        for words in para:
-          text += " ".join(words)
-      textsToExtract[l] = text
-  # delete folder after finishing getting text
-  cmd = f'rm -r data/{id}'
-  result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-  return textsToExtract
-# Step 3.2: Extract context
-def extract_context(text, keyword, window=500):
-    # firstly try accession number
-    idx = text.find(keyword)
-    if idx == -1:
-        return "Sample ID not found."
-    return text[max(0, idx-window): idx+window]
-def extract_relevant_paragraphs(text, accession, keep_if=None, isolate=None):
-    if keep_if is None:
-        keep_if = ["sample", "method", "mtdna", "sequence", "collected", "dataset", "supplementary", "table"]
-
-    outputs = ""
-    text = text.lower()
-
-    # If isolate is provided, prioritize paragraphs that mention it
-    # If isolate is provided, prioritize paragraphs that mention it
-    if accession and accession.lower() in text:
-        if extract_context(text, accession.lower(), window=700) != "Sample ID not found.":
-            outputs += extract_context(text, accession.lower(), window=700)       
-    if isolate and isolate.lower() in text:
-        if extract_context(text, isolate.lower(), window=700) != "Sample ID not found.":
-            outputs += extract_context(text, isolate.lower(), window=700)
-    for keyword in keep_if:
-        para = extract_context(text, keyword)
-        if para and para not in outputs:
-            outputs += para + "\n"
-    return outputs
-# Step 4: Classification for now (demo purposes)
-# 4.1: Using a HuggingFace model (question-answering)
-def infer_fromQAModel(context, question="Where is the mtDNA sample from?"):
-    try:
-        qa = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
-        result = qa({"context": context, "question": question})
-        return result.get("answer", "Unknown")
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# 4.2: Infer from haplogroup
-# Load pre-trained spaCy model for NER
-try:
-    nlp = spacy.load("en_core_web_sm")
-except OSError:
-    download("en_core_web_sm")
-    nlp = spacy.load("en_core_web_sm")
-
-# Define the haplogroup-to-region mapping (simple rule-based)
-import csv
-
-def load_haplogroup_mapping(csv_path):
-    mapping = {}
-    with open(csv_path) as f:
-        reader = csv.DictReader(f)
-        for row in reader:
-            mapping[row["haplogroup"]] = [row["region"],row["source"]]
-    return mapping
-
-# Function to extract haplogroup from the text
-def extract_haplogroup(text):
-    match = re.search(r'\bhaplogroup\s+([A-Z][0-9a-z]*)\b', text)
-    if match:
-        submatch = re.match(r'^[A-Z][0-9]*', match.group(1))
-        if submatch:
-            return submatch.group(0)
-        else:
-            return match.group(1)  # fallback
-    fallback = re.search(r'\b([A-Z][0-9a-z]{1,5})\b', text)
-    if fallback:
-        return fallback.group(1)
-    return None
-
-
-# Function to extract location based on NER
-def extract_location(text):
-    doc = nlp(text)
-    locations = []
-    for ent in doc.ents:
-        if ent.label_ == "GPE":  # GPE = Geopolitical Entity (location)
-            locations.append(ent.text)
-    return locations
-
-# Function to infer location from haplogroup
-def infer_location_from_haplogroup(haplogroup):
-  haplo_map = load_haplogroup_mapping("data/haplogroup_regions_extended.csv")
-  return haplo_map.get(haplogroup, ["Unknown","Unknown"])
-
-# Function to classify the mtDNA sample
-def classify_mtDNA_sample_from_haplo(text):
-    # Extract haplogroup
-    haplogroup = extract_haplogroup(text)
-    # Extract location based on NER
-    locations = extract_location(text)
-    # Infer location based on haplogroup
-    inferred_location, sourceHaplo = infer_location_from_haplogroup(haplogroup)[0],infer_location_from_haplogroup(haplogroup)[1]
-    return {
-        "source":sourceHaplo,
-        "locations_found_in_context": locations,
-        "haplogroup": haplogroup,
-        "inferred_location": inferred_location
-
-    }
-# 4.3 Get from available NCBI
-def infer_location_fromNCBI(accession):
-    try:
-        handle = Entrez.efetch(db="nuccore", id=accession, rettype="medline", retmode="text")
-        text = handle.read()
-        handle.close()
-        match = re.search(r'/(geo_loc_name|country|location)\s*=\s*"([^"]+)"', text)
-        if match:
-            return match.group(2), match.group(0)  # This is the value like "Brunei"
-        return "Not found", "Not found"
-
-    except Exception as e:
-        print("❌ Entrez error:", e)
-        return "Not found", "Not found"
-
-### ANCIENT/MODERN FLAG
-from Bio import Entrez
-import re
-
-def flag_ancient_modern(accession, textsToExtract, isolate=None):
-    """
-    Try to classify a sample as Ancient or Modern using:
-    1. NCBI accession (if available)
-    2. Supplementary text or context fallback
-    """
-    context = ""
-    label, explain = "", ""
-
-    try:
-        # Check if we can fetch metadata from NCBI using the accession
-        handle = Entrez.efetch(db="nuccore", id=accession, rettype="medline", retmode="text")
-        text = handle.read()
-        handle.close()
-
-        isolate_source = re.search(r'/(isolation_source)\s*=\s*"([^"]+)"', text)
-        if isolate_source:
-            context += isolate_source.group(0) + " "
-
-        specimen = re.search(r'/(specimen|specimen_voucher)\s*=\s*"([^"]+)"', text)
-        if specimen:
-            context += specimen.group(0) + " "
-
-        if context.strip():
-            label, explain = detect_ancient_flag(context)
-            if label!="Unknown":
-              return label, explain + " from NCBI\n(" + context + ")"
-
-        # If no useful NCBI metadata, check supplementary texts
-        if textsToExtract:
-            labels = {"modern": [0, ""], "ancient": [0, ""], "unknown": 0}
-
-            for source in textsToExtract:
-                text_block = textsToExtract[source]
-                context = extract_relevant_paragraphs(text_block, accession, isolate=isolate)  # Reduce to informative paragraph(s)
-                label, explain = detect_ancient_flag(context)
-
-                if label == "Ancient":
-                    labels["ancient"][0] += 1
-                    labels["ancient"][1] += f"{source}:\n{explain}\n\n"
-                elif label == "Modern":
-                    labels["modern"][0] += 1
-                    labels["modern"][1] += f"{source}:\n{explain}\n\n"
-                else:
-                    labels["unknown"] += 1
-
-            if max(labels["modern"][0],labels["ancient"][0]) > 0:
-                if labels["modern"][0] > labels["ancient"][0]:
-                    return "Modern", labels["modern"][1]
-                else:
-                    return "Ancient", labels["ancient"][1]
-            else:
-              return "Unknown", "No strong keywords detected"
-        else:
-            print("No DOI or PubMed ID available for inference.")
-            return "", ""
-
-    except Exception as e:
-        print("Error:", e)
-        return "", ""
-
-
-def detect_ancient_flag(context_snippet):
-    context = context_snippet.lower()
-
-    ancient_keywords = [
-        "ancient", "archaeological", "prehistoric", "neolithic", "mesolithic", "paleolithic",
-        "bronze age", "iron age", "burial", "tomb", "skeleton", "14c", "radiocarbon", "carbon dating",
-        "postmortem damage", "udg treatment", "adna", "degradation", "site", "excavation",
-        "archaeological context", "temporal transect", "population replacement", "cal bp", "calbp", "carbon dated"
-    ]
-
-    modern_keywords = [
-        "modern", "hospital", "clinical", "consent","blood","buccal","unrelated", "blood sample","buccal sample","informed consent", "donor", "healthy", "patient",
-        "genotyping", "screening", "medical", "cohort", "sequencing facility", "ethics approval",
-        "we analysed", "we analyzed", "dataset includes", "new sequences", "published data",
-        "control cohort", "sink population", "genbank accession", "sequenced", "pipeline", 
-        "bioinformatic analysis", "samples from", "population genetics", "genome-wide data"
-    ]
-
-    ancient_hits = [k for k in ancient_keywords if k in context]
-    modern_hits = [k for k in modern_keywords if k in context]
-
-    if ancient_hits and not modern_hits:
-        return "Ancient", f"Flagged as ancient due to keywords: {', '.join(ancient_hits)}"
-    elif modern_hits and not ancient_hits:
-        return "Modern", f"Flagged as modern due to keywords: {', '.join(modern_hits)}"
-    elif ancient_hits and modern_hits:
-        if len(ancient_hits) >= len(modern_hits):
-            return "Ancient", f"Mixed context, leaning ancient due to: {', '.join(ancient_hits)}"
-        else:
-            return "Modern", f"Mixed context, leaning modern due to: {', '.join(modern_hits)}"
-    
-    # Fallback to QA
-    answer = infer_fromQAModel(context, question="Are the mtDNA samples ancient or modern? Explain why.")
-    if answer.startswith("Error"):
-        return "Unknown", answer
-    if "ancient" in answer.lower():
-        return "Ancient", f"Leaning ancient based on QA: {answer}"
-    elif "modern" in answer.lower():
-        return "Modern", f"Leaning modern based on QA: {answer}"
-    else:
-        return "Unknown", f"No strong keywords or QA clues. QA said: {answer}"
-
-# STEP 5: Main pipeline: accession -> 1. get pubmed id and isolate -> 2. get doi -> 3. get text -> 4. prediction -> 5. output: inferred location + explanation + confidence score
-def classify_sample_location(accession):
-  outputs = {}
-  keyword, context, location, qa_result, haplo_result = "", "", "", "", ""
-  # Step 1: get pubmed id and isolate
-  pubmedID, isolate = get_info_from_accession(accession)
-  '''if not pubmedID:
-    return {"error": f"Could not retrieve PubMed ID for accession {accession}"}'''
-  if not isolate:
-    isolate = "UNKNOWN_ISOLATE"
-  # Step 2: get doi
-  doi = get_doi_from_pubmed_id(pubmedID)
-  '''if not doi:
-    return {"error": "DOI not found for this accession. Cannot fetch paper or context."}'''
-  # Step 3: get text
-  '''textsToExtract = { "doiLink":"paperText"
-                        "file1.pdf":"text1",
-                        "file2.doc":"text2",
-                        "file3.xlsx":excelText3'''
-  if doi and pubmedID:                      
-    textsToExtract = get_paper_text(doi,pubmedID)
-  else: textsToExtract = {}  
-  '''if not textsToExtract:
-    return {"error": f"No texts extracted for DOI {doi}"}'''
-  if isolate not in [None, "UNKNOWN_ISOLATE"]:
-    label, explain = flag_ancient_modern(accession,textsToExtract,isolate)
-  else: 
-    label, explain = flag_ancient_modern(accession,textsToExtract)  
-  # Step 4: prediction
-  outputs[accession] = {}
-  outputs[isolate] = {}
-  # 4.0 Infer from NCBI
-  location, outputNCBI = infer_location_fromNCBI(accession)
-  NCBI_result = {
-      "source": "NCBI",
-      "sample_id": accession,
-      "predicted_location": location,
-      "context_snippet": outputNCBI}
-  outputs[accession]["NCBI"]= {"NCBI": NCBI_result}
-  if textsToExtract:
-    long_text = ""
-    for key in textsToExtract:
-      text = textsToExtract[key]
-      # try accession number first
-      outputs[accession][key] = {}
-      keyword = accession
-      context = extract_context(text, keyword, window=500)
-      # 4.1: Using a HuggingFace model (question-answering)
-      location = infer_fromQAModel(context, question=f"Where is the mtDNA sample {keyword} from?")
-      qa_result = {
-          "source": key,
-          "sample_id": keyword,
-          "predicted_location": location,
-          "context_snippet": context
-      }
-      outputs[keyword][key]["QAModel"] = qa_result
-      # 4.2: Infer from haplogroup
-      haplo_result = classify_mtDNA_sample_from_haplo(context)
-      outputs[keyword][key]["haplogroup"] = haplo_result
-      # try isolate
-      keyword = isolate
-      outputs[isolate][key] = {}
-      context = extract_context(text, keyword, window=500)
-      # 4.1.1: Using a HuggingFace model (question-answering)
-      location = infer_fromQAModel(context, question=f"Where is the mtDNA sample {keyword} from?")
-      qa_result = {
-          "source": key,
-          "sample_id": keyword,
-          "predicted_location": location,
-          "context_snippet": context
-      }
-      outputs[keyword][key]["QAModel"] = qa_result
-      # 4.2.1: Infer from haplogroup
-      haplo_result = classify_mtDNA_sample_from_haplo(context)
-      outputs[keyword][key]["haplogroup"] = haplo_result
-      # add long text
-      long_text += text + ". \n"
-    # 4.3: UpgradeClassify
-    # try sample_id as accession number
-    sample_id = accession
-    if sample_id:
-      filtered_context = filter_context_for_sample(sample_id.upper(), long_text, window_size=1)
-      locations = infer_location_for_sample(sample_id.upper(), filtered_context)
-      if locations!="No clear location found in top matches":
-        outputs[sample_id]["upgradeClassifier"] = {}
-        outputs[sample_id]["upgradeClassifier"]["upgradeClassifier"] = {
-          "source": "From these sources combined: "+ ", ".join(list(textsToExtract.keys())),
-          "sample_id": sample_id,
-          "predicted_location": ", ".join(locations),
-          "context_snippep": "First 1000 words: \n"+ filtered_context[:1000]
-        }
-    # try sample_id as isolate name
-    sample_id = isolate
-    if sample_id:
-      filtered_context = filter_context_for_sample(sample_id.upper(), long_text, window_size=1)
-      locations = infer_location_for_sample(sample_id.upper(), filtered_context)
-      if locations!="No clear location found in top matches":
-        outputs[sample_id]["upgradeClassifier"] = {}
-        outputs[sample_id]["upgradeClassifier"]["upgradeClassifier"] = {
-          "source": "From these sources combined: "+ ", ".join(list(textsToExtract.keys())),
-          "sample_id": sample_id,
-          "predicted_location": ", ".join(locations),
-          "context_snippep": "First 1000 words: \n"+ filtered_context[:1000]
-        }
-  return outputs, label, explain
+# mtDNA Location Classifier MVP (Google Colab)
+# Accepts accession number → Fetches PubMed ID + isolate name → Gets abstract → Predicts location
+import os
+#import streamlit as st
+import subprocess
+import re
+from Bio import Entrez
+import fitz
+import spacy
+from spacy.cli import download
+from NER.PDF import pdf
+from NER.WordDoc import wordDoc
+from NER.html import extractHTML
+from NER.word2Vec import word2vec
+from transformers import pipeline
+import urllib.parse, requests
+from pathlib import Path
+from upgradeClassify import filter_context_for_sample, infer_location_for_sample
+
+# Set your email (required by NCBI Entrez)
+#Entrez.email = "[email protected]"
+import nltk
+
+nltk.download("stopwords")
+nltk.download("punkt")
+nltk.download('punkt_tab')
+# Step 1: Get PubMed ID from Accession using EDirect
+from Bio import Entrez, Medline
+import re
+
+Entrez.email = "your_email@example.com"
+
+# --- Helper Functions (Re-organized and Upgraded) ---
+
+def fetch_ncbi_metadata(accession_number):
+    """
+    Fetches metadata directly from NCBI GenBank using Entrez.
+    Includes robust error handling and improved field extraction.
+    Prioritizes location extraction from geo_loc_name, then notes, then other qualifiers.
+    Also attempts to extract ethnicity and sample_type (ancient/modern).
+
+    Args:
+        accession_number (str): The NCBI accession number (e.g., "ON792208").
+
+    Returns:
+        dict: A dictionary containing 'country', 'specific_location', 'ethnicity',
+              'sample_type', 'collection_date', 'isolate', 'title', 'doi', 'pubmed_id'.
+    """
+    Entrez.email = "[email protected]" # Required by NCBI, REPLACE WITH YOUR EMAIL
+
+    country = "unknown"
+    specific_location = "unknown"
+    ethnicity = "unknown"
+    sample_type = "unknown"
+    collection_date = "unknown"
+    isolate = "unknown"
+    title = "unknown"
+    doi = "unknown"
+    pubmed_id = None
+    all_feature = "unknown"
+
+    KNOWN_COUNTRIES = [
+        "Afghanistan", "Albania", "Algeria", "Andorra", "Angola", "Antigua and Barbuda", "Argentina", "Armenia", "Australia", "Austria", "Azerbaijan",
+        "Bahamas", "Bahrain", "Bangladesh", "Barbados", "Belarus", "Belgium", "Belize", "Benin", "Bhutan", "Bolivia", "Bosnia and Herzegovina", "Botswana", "Brazil", "Brunei", "Bulgaria", "Burkina Faso", "Burundi",
+        "Cabo Verde", "Cambodia", "Cameroon", "Canada", "Central African Republic", "Chad", "Chile", "China", "Colombia", "Comoros", "Congo (Brazzaville)", "Congo (Kinshasa)", "Costa Rica", "Croatia", "Cuba", "Cyprus", "Czechia",
+        "Denmark", "Djibouti", "Dominica", "Dominican Republic", "Ecuador", "Egypt", "El Salvador", "Equatorial Guinea", "Eritrea", "Estonia", "Eswatini", "Ethiopia",
+        "Fiji", "Finland", "France", "Gabon", "Gambia", "Georgia", "Germany", "Ghana", "Greece", "Grenada", "Guatemala", "Guinea", "Guinea-Bissau", "Guyana",
+        "Haiti", "Honduras", "Hungary", "Iceland", "India", "Indonesia", "Iran", "Iraq", "Ireland", "Israel", "Italy", "Ivory Coast", "Jamaica", "Japan", "Jordan",
+        "Kazakhstan", "Kenya", "Kiribati", "Kosovo", "Kuwait", "Kyrgyzstan", "Laos", "Latvia", "Lebanon", "Lesotho", "Liberia", "Libya", "Liechtenstein", "Lithuania", "Luxembourg",
+        "Madagascar", "Malawi", "Malaysia", "Maldives", "Mali", "Malta", "Marshall Islands", "Mauritania", "Mauritius", "Mexico", "Micronesia", "Moldova", "Monaco", "Mongolia", "Montenegro", "Morocco", "Mozambique", "Myanmar",
+        "Namibia", "Nauru", "Nepal", "Netherlands", "New Zealand", "Nicaragua", "Niger", "Nigeria", "North Korea", "North Macedonia", "Norway", "Oman",
+        "Pakistan", "Palau", "Palestine", "Panama", "Papua New Guinea", "Paraguay", "Peru", "Philippines", "Poland", "Portugal", "Qatar", "Romania", "Russia", "Rwanda",
+        "Saint Kitts and Nevis", "Saint Lucia", "Saint Vincent and the Grenadines", "Samoa", "San Marino", "Sao Tome and Principe", "Saudi Arabia", "Senegal", "Serbia", "Seychelles", "Sierra Leone", "Singapore", "Slovakia", "Slovenia", "Solomon Islands", "Somalia", "South Africa", "South Korea", "South Sudan", "Spain", "Sri Lanka", "Sudan", "Suriname", "Sweden", "Switzerland", "Syria",
+        "Taiwan", "Tajikistan", "Tanzania", "Thailand", "Timor-Leste", "Togo", "Tonga", "Trinidad and Tobago", "Tunisia", "Turkey", "Turkmenistan", "Tuvalu",
+        "Uganda", "Ukraine", "United Arab Emirates", "United Kingdom", "United States", "Uruguay", "Uzbekistan", "Vanuatu", "Vatican City", "Venezuela", "Vietnam",
+        "Yemen", "Zambia", "Zimbabwe"
+    ]
+    COUNTRY_PATTERN = re.compile(r'\b(' + '|'.join(re.escape(c) for c in KNOWN_COUNTRIES) + r')\b', re.IGNORECASE)
+
+    try:
+        handle = Entrez.efetch(db="nucleotide", id=str(accession_number), rettype="gb", retmode="xml")
+        record = Entrez.read(handle)
+        handle.close()
+
+        gb_seq = None
+        # Validate record structure: It should be a list with at least one element (a dict)
+        if isinstance(record, list) and len(record) > 0:
+            if isinstance(record[0], dict):
+                gb_seq = record[0]
+            else:
+                print(f"Warning: record[0] is not a dictionary for {accession_number}. Type: {type(record[0])}")
+        else:
+            print(f"Warning: No valid record or empty record list from NCBI for {accession_number}.")
+
+        # If gb_seq is still None, return defaults
+        if gb_seq is None:
+            return {"country": "unknown", "specific_location": "unknown", "ethnicity": "unknown",
+                    "sample_type": "unknown", "collection_date": "unknown", "isolate": "unknown",
+                    "title": "unknown", "doi": "unknown", "pubmed_id": None}
+
+
+        # If gb_seq is valid, proceed with extraction
+        collection_date = gb_seq.get("GBSeq_create-date","unknown")
+
+        references = gb_seq.get("GBSeq_references", [])
+        for ref in references:
+            if not pubmed_id:
+                pubmed_id = ref.get("GBReference_pubmed",None)
+            if title == "unknown":
+                title = ref.get("GBReference_title","unknown")
+            for xref in ref.get("GBReference_xref", []):
+                if xref.get("GBXref_dbname") == "doi":
+                    doi = xref.get("GBXref_id")
+                    break
+
+        features = gb_seq.get("GBSeq_feature-table", [])
+
+        context_for_flagging = "" # Accumulate text for ancient/modern detection
+        features_context = ""
+        for feature in features:
+            if feature.get("GBFeature_key") == "source":
+                feature_context = ""
+                qualifiers = feature.get("GBFeature_quals", [])
+                found_country = "unknown"
+                found_specific_location = "unknown"
+                found_ethnicity = "unknown"
+
+                temp_geo_loc_name = "unknown"
+                temp_note_origin_locality = "unknown"
+                temp_country_qual = "unknown"
+                temp_locality_qual = "unknown"
+                temp_collection_location_qual = "unknown"
+                temp_isolation_source_qual = "unknown"
+                temp_env_sample_qual = "unknown"
+                temp_pop_qual = "unknown"
+                temp_organism_qual = "unknown"
+                temp_specimen_qual = "unknown"
+                temp_strain_qual = "unknown"
+
+                for qual in qualifiers:
+                    qual_name = qual.get("GBQualifier_name")
+                    qual_value = qual.get("GBQualifier_value")
+                    feature_context += qual_name + ": " + qual_value +"\n"
+                    if qual_name == "collection_date":
+                        collection_date = qual_value
+                    elif qual_name == "isolate":
+                        isolate = qual_value
+                    elif qual_name == "population":
+                        temp_pop_qual = qual_value
+                    elif qual_name == "organism":
+                        temp_organism_qual = qual_value
+                    elif qual_name == "specimen_voucher" or qual_name == "specimen":
+                        temp_specimen_qual = qual_value
+                    elif qual_name == "strain":
+                        temp_strain_qual = qual_value
+                    elif qual_name == "isolation_source":
+                        temp_isolation_source_qual = qual_value
+                    elif qual_name == "environmental_sample":
+                        temp_env_sample_qual = qual_value
+
+                    if qual_name == "geo_loc_name": temp_geo_loc_name = qual_value
+                    elif qual_name == "note":
+                        if qual_value.startswith("origin_locality:"):
+                            temp_note_origin_locality = qual_value
+                        context_for_flagging += qual_value + " " # Capture all notes for flagging  
+                    elif qual_name == "country": temp_country_qual = qual_value
+                    elif qual_name == "locality": temp_locality_qual = qual_value
+                    elif qual_name == "collection_location": temp_collection_location_qual = qual_value
+
+
+                # --- Aggregate all relevant info into context_for_flagging ---
+                context_for_flagging += f" {isolate} {temp_isolation_source_qual} {temp_specimen_qual} {temp_strain_qual} {temp_organism_qual} {temp_geo_loc_name} {temp_collection_location_qual} {temp_env_sample_qual}"
+                context_for_flagging = context_for_flagging.strip()
+                
+                # --- Determine final country and specific_location based on priority ---
+                if temp_geo_loc_name != "unknown":
+                    parts = [p.strip() for p in temp_geo_loc_name.split(':')]
+                    if len(parts) > 1: 
+                      found_specific_location = parts[-1]; found_country = parts[0]
+                    else: found_country = temp_geo_loc_name; found_specific_location = "unknown"
+                elif temp_note_origin_locality != "unknown":
+                    match = re.search(r"origin_locality:\s*(.*)", temp_note_origin_locality, re.IGNORECASE)
+                    if match:
+                        location_string = match.group(1).strip()
+                        parts = [p.strip() for p in location_string.split(':')]
+                        if len(parts) > 1: found_country = parts[-1]; found_specific_location = parts[0]
+                        else: found_country = location_string; found_specific_location = "unknown"
+                elif temp_locality_qual != "unknown":
+                    found_country_match = COUNTRY_PATTERN.search(temp_locality_qual)
+                    if found_country_match: found_country = found_country_match.group(1); temp_loc = re.sub(re.escape(found_country), '', temp_locality_qual, flags=re.IGNORECASE).strip().replace(',', '').replace(':', '').replace(';', '').strip(); found_specific_location = temp_loc if temp_loc else "unknown"
+                    else: found_specific_location = temp_locality_qual; found_country = "unknown"
+                elif temp_collection_location_qual != "unknown":
+                    found_country_match = COUNTRY_PATTERN.search(temp_collection_location_qual)
+                    if found_country_match: found_country = found_country_match.group(1); temp_loc = re.sub(re.escape(found_country), '', temp_collection_location_qual, flags=re.IGNORECASE).strip().replace(',', '').replace(':', '').replace(';', '').strip(); found_specific_location = temp_loc if temp_loc else "unknown"
+                    else: found_specific_location = temp_collection_location_qual; found_country = "unknown"
+                elif temp_isolation_source_qual != "unknown":
+                    found_country_match = COUNTRY_PATTERN.search(temp_isolation_source_qual)
+                    if found_country_match: found_country = found_country_match.group(1); temp_loc = re.sub(re.escape(found_country), '', temp_isolation_source_qual, flags=re.IGNORECASE).strip().replace(',', '').replace(':', '').replace(';', '').strip(); found_specific_location = temp_loc if temp_loc else "unknown"
+                    else: found_specific_location = temp_isolation_source_qual; found_country = "unknown"
+                elif temp_env_sample_qual != "unknown":
+                    found_country_match = COUNTRY_PATTERN.search(temp_env_sample_qual)
+                    if found_country_match: found_country = found_country_match.group(1); temp_loc = re.sub(re.escape(found_country), '', temp_env_sample_qual, flags=re.IGNORECASE).strip().replace(',', '').replace(':', '').replace(';', '').strip(); found_specific_location = temp_loc if temp_loc else "unknown"
+                    else: found_specific_location = temp_env_sample_qual; found_country = "unknown"
+                if found_country == "unknown" and temp_country_qual != "unknown":
+                     found_country_match = COUNTRY_PATTERN.search(temp_country_qual)
+                     if found_country_match: found_country = found_country_match.group(1)
+
+                country = found_country
+                specific_location = found_specific_location
+                # --- Determine final ethnicity ---
+                if temp_pop_qual != "unknown":
+                    found_ethnicity = temp_pop_qual
+                elif isolate != "unknown" and re.fullmatch(r'[A-Za-z\s\-]+', isolate) and get_country_from_text(isolate) == "unknown":
+                     found_ethnicity = isolate
+                elif context_for_flagging != "unknown": # Use the broader context for ethnicity patterns
+                    eth_match = re.search(r'(?:population|ethnicity|isolate source):\s*([A-Za-z\s\-]+)', context_for_flagging, re.IGNORECASE)
+                    if eth_match:
+                        found_ethnicity = eth_match.group(1).strip()
+
+                ethnicity = found_ethnicity
+
+                # --- Determine sample_type (ancient/modern) ---
+                if context_for_flagging:
+                    sample_type, explain = detect_ancient_flag(context_for_flagging)
+                features_context += feature_context + "\n"
+                break
+
+        if specific_location != "unknown" and specific_location.lower() == country.lower():
+            specific_location = "unknown"
+        if not features_context:  features_context = "unknown"    
+        return {"country": country.lower(),
+                "specific_location": specific_location.lower(),
+                "ethnicity": ethnicity.lower(),
+                "sample_type": sample_type.lower(),
+                "collection_date": collection_date,
+                "isolate": isolate,
+                "title": title,
+                "doi": doi,
+                "pubmed_id": pubmed_id,
+                "all_features": features_context}
+
+    except Exception as e:
+        print(f"Error fetching NCBI data for {accession_number}: {e}")
+        return {"country": "unknown",
+                "specific_location": "unknown",
+                "ethnicity": "unknown",
+                "sample_type": "unknown",
+                "collection_date": "unknown",
+                "isolate": "unknown",
+                "title": "unknown",
+                "doi": "unknown",
+                "pubmed_id": None,
+                "all_features": "unknown"}
+
+# --- Helper function for country matching (re-defined from main code to be self-contained) ---
+_country_keywords = {
+    "thailand": "Thailand", "laos": "Laos", "cambodia": "Cambodia", "myanmar": "Myanmar",
+    "philippines": "Philippines", "indonesia": "Indonesia", "malaysia": "Malaysia",
+    "china": "China", "chinese": "China", "india": "India", "taiwan": "Taiwan",
+    "vietnam": "Vietnam", "russia": "Russia", "siberia": "Russia", "nepal": "Nepal",
+    "japan": "Japan", "sumatra": "Indonesia", "borneu": "Indonesia",
+    "yunnan": "China", "tibet": "China", "northern mindanao": "Philippines",
+    "west malaysia": "Malaysia", "north thailand": "Thailand", "central thailand": "Thailand",
+    "northeast thailand": "Thailand", "east myanmar": "Myanmar", "west thailand": "Thailand",
+    "central india": "India", "east india": "India", "northeast india": "India",
+    "south sibera": "Russia", "mongolia": "China", "beijing": "China", "south korea": "South Korea",
+    "north asia": "unknown", "southeast asia": "unknown", "east asia": "unknown"
+}
+
+def get_country_from_text(text):
+    text_lower = text.lower()
+    for keyword, country in _country_keywords.items():
+        if keyword in text_lower:
+            return country
+    return "unknown"
+# The result will be seen as manualLink for the function get_paper_text
+def search_google_custom(query, max_results=3):
+  # query should be the title from ncbi or paper/source title
+    GOOGLE_CSE_API_KEY = "AIzaSyAg_Hi5DPit2bvvwCs1PpUkAPRZun7yCRQ"
+    GOOGLE_CSE_CX = "25a51c433f148490c"
+    endpoint = "https://www.googleapis.com/customsearch/v1"
+    params = {
+        "key": GOOGLE_CSE_API_KEY,
+        "cx": GOOGLE_CSE_CX,
+        "q": query,
+        "num": max_results
+    }
+    try:
+        response = requests.get(endpoint, params=params)
+        if response.status_code == 429:
+            print("Rate limit hit. Try again later.")
+            return []
+        response.raise_for_status()
+        data = response.json().get("items", [])
+        return [item.get("link") for item in data if item.get("link")]
+    except Exception as e:
+        print("Google CSE error:", e)
+        return []
+# Step 3: Extract Text: Get the paper (html text), sup. materials (pdf, doc, excel) and do text-preprocessing
+# Step 3.1: Extract Text
+# sub: download excel file
+def download_excel_file(url, save_path="temp.xlsx"):
+    if "view.officeapps.live.com" in url:
+        parsed_url = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
+        real_url = urllib.parse.unquote(parsed_url["src"][0])
+        response = requests.get(real_url)
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+        return save_path
+    elif url.startswith("http") and (url.endswith(".xls") or url.endswith(".xlsx")):
+        response = requests.get(url)
+        response.raise_for_status()  # Raises error if download fails
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+        return save_path
+    else:
+        print("URL must point directly to an .xls or .xlsx file\n or it already downloaded.")
+        return url
+def get_paper_text(doi,id,manualLinks=None):
+  # create the temporary folder to contain the texts
+  folder_path = Path("data/"+str(id))
+  if not folder_path.exists():
+      cmd = f'mkdir data/{id}'
+      result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+      print("data/"+str(id) +" created.")
+  else:
+      print("data/"+str(id) +" already exists.")
+  saveLinkFolder = "data/"+id
+
+  link = 'https://doi.org/' + doi
+  '''textsToExtract = { "doiLink":"paperText"
+                        "file1.pdf":"text1",
+                        "file2.doc":"text2",
+                        "file3.xlsx":excelText3'''
+  textsToExtract = {}
+  # get the file to create listOfFile for each id
+  html = extractHTML.HTML("",link)
+  jsonSM = html.getSupMaterial()
+  text = ""
+  links  = [link] + sum((jsonSM[key] for key in jsonSM),[])
+  if manualLinks != None:
+    links += manualLinks
+  for l in links:
+    # get the main paper
+    name = l.split("/")[-1]
+    file_path = folder_path / name
+    if l == link:
+      text = html.getListSection()
+      textsToExtract[link] = text
+    elif l.endswith(".pdf"):
+      if file_path.is_file():
+          l = saveLinkFolder + "/" + name
+          print("File exists.")
+      p = pdf.PDF(l,saveLinkFolder,doi)
+      f = p.openPDFFile()
+      pdf_path = saveLinkFolder + "/" + l.split("/")[-1]
+      doc = fitz.open(pdf_path)
+      text = "\n".join([page.get_text() for page in doc])
+      textsToExtract[l] = text
+    elif l.endswith(".doc") or l.endswith(".docx"):
+      d = wordDoc.wordDoc(l,saveLinkFolder)
+      text = d.extractTextByPage()
+      textsToExtract[l] = text
+    elif l.split(".")[-1].lower() in "xlsx":
+      wc = word2vec.word2Vec()
+      # download excel file if it not downloaded yet
+      savePath = saveLinkFolder +"/"+ l.split("/")[-1]
+      excelPath = download_excel_file(l, savePath)
+      corpus = wc.tableTransformToCorpusText([],excelPath)
+      text = ''
+      for c in corpus:
+        para = corpus[c]
+        for words in para:
+          text += " ".join(words)
+      textsToExtract[l] = text
+  # delete folder after finishing getting text
+  #cmd = f'rm -r data/{id}'
+  #result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+  return textsToExtract
+# Step 3.2: Extract context
+def extract_context(text, keyword, window=500):
+    # firstly try accession number
+    idx = text.find(keyword)
+    if idx == -1:
+        return "Sample ID not found."
+    return text[max(0, idx-window): idx+window]
+def extract_relevant_paragraphs(text, accession, keep_if=None, isolate=None):
+    if keep_if is None:
+        keep_if = ["sample", "method", "mtdna", "sequence", "collected", "dataset", "supplementary", "table"]
+
+    outputs = ""
+    text = text.lower()
+
+    # If isolate is provided, prioritize paragraphs that mention it
+    # If isolate is provided, prioritize paragraphs that mention it
+    if accession and accession.lower() in text:
+        if extract_context(text, accession.lower(), window=700) != "Sample ID not found.":
+            outputs += extract_context(text, accession.lower(), window=700)       
+    if isolate and isolate.lower() in text:
+        if extract_context(text, isolate.lower(), window=700) != "Sample ID not found.":
+            outputs += extract_context(text, isolate.lower(), window=700)
+    for keyword in keep_if:
+        para = extract_context(text, keyword)
+        if para and para not in outputs:
+            outputs += para + "\n"
+    return outputs
+# Step 4: Classification for now (demo purposes)
+# 4.1: Using a HuggingFace model (question-answering)
+def infer_fromQAModel(context, question="Where is the mtDNA sample from?"):
+    try:
+        qa = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
+        result = qa({"context": context, "question": question})
+        return result.get("answer", "Unknown")
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# 4.2: Infer from haplogroup
+# Load pre-trained spaCy model for NER
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    download("en_core_web_sm")
+    nlp = spacy.load("en_core_web_sm")
+
+# Define the haplogroup-to-region mapping (simple rule-based)
+import csv
+
+def load_haplogroup_mapping(csv_path):
+    mapping = {}
+    with open(csv_path) as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            mapping[row["haplogroup"]] = [row["region"],row["source"]]
+    return mapping
+
+# Function to extract haplogroup from the text
+def extract_haplogroup(text):
+    match = re.search(r'\bhaplogroup\s+([A-Z][0-9a-z]*)\b', text)
+    if match:
+        submatch = re.match(r'^[A-Z][0-9]*', match.group(1))
+        if submatch:
+            return submatch.group(0)
+        else:
+            return match.group(1)  # fallback
+    fallback = re.search(r'\b([A-Z][0-9a-z]{1,5})\b', text)
+    if fallback:
+        return fallback.group(1)
+    return None
+
+
+# Function to extract location based on NER
+def extract_location(text):
+    doc = nlp(text)
+    locations = []
+    for ent in doc.ents:
+        if ent.label_ == "GPE":  # GPE = Geopolitical Entity (location)
+            locations.append(ent.text)
+    return locations
+
+# Function to infer location from haplogroup
+def infer_location_from_haplogroup(haplogroup):
+  haplo_map = load_haplogroup_mapping("data/haplogroup_regions_extended.csv")
+  return haplo_map.get(haplogroup, ["Unknown","Unknown"])
+
+# Function to classify the mtDNA sample
+def classify_mtDNA_sample_from_haplo(text):
+    # Extract haplogroup
+    haplogroup = extract_haplogroup(text)
+    # Extract location based on NER
+    locations = extract_location(text)
+    # Infer location based on haplogroup
+    inferred_location, sourceHaplo = infer_location_from_haplogroup(haplogroup)[0],infer_location_from_haplogroup(haplogroup)[1]
+    return {
+        "source":sourceHaplo,
+        "locations_found_in_context": locations,
+        "haplogroup": haplogroup,
+        "inferred_location": inferred_location
+
+    }
+# 4.3 Get from available NCBI
+def infer_location_fromNCBI(accession):
+    try:
+        handle = Entrez.efetch(db="nuccore", id=accession, rettype="medline", retmode="text")
+        text = handle.read()
+        handle.close()
+        match = re.search(r'/(geo_loc_name|country|location)\s*=\s*"([^"]+)"', text)
+        if match:
+            return match.group(2), match.group(0)  # This is the value like "Brunei"
+        return "Not found", "Not found"
+
+    except Exception as e:
+        print("❌ Entrez error:", e)
+        return "Not found", "Not found"
+
+### ANCIENT/MODERN FLAG
+from Bio import Entrez
+import re
+
+def flag_ancient_modern(accession, textsToExtract, isolate=None):
+    """
+    Try to classify a sample as Ancient or Modern using:
+    1. NCBI accession (if available)
+    2. Supplementary text or context fallback
+    """
+    context = ""
+    label, explain = "", ""
+
+    try:
+        # Check if we can fetch metadata from NCBI using the accession
+        handle = Entrez.efetch(db="nuccore", id=accession, rettype="medline", retmode="text")
+        text = handle.read()
+        handle.close()
+
+        isolate_source = re.search(r'/(isolation_source)\s*=\s*"([^"]+)"', text)
+        if isolate_source:
+            context += isolate_source.group(0) + " "
+
+        specimen = re.search(r'/(specimen|specimen_voucher)\s*=\s*"([^"]+)"', text)
+        if specimen:
+            context += specimen.group(0) + " "
+
+        if context.strip():
+            label, explain = detect_ancient_flag(context)
+            if label!="Unknown":
+              return label, explain + " from NCBI\n(" + context + ")"
+
+        # If no useful NCBI metadata, check supplementary texts
+        if textsToExtract:
+            labels = {"modern": [0, ""], "ancient": [0, ""], "unknown": 0}
+
+            for source in textsToExtract:
+                text_block = textsToExtract[source]
+                context = extract_relevant_paragraphs(text_block, accession, isolate=isolate)  # Reduce to informative paragraph(s)
+                label, explain = detect_ancient_flag(context)
+
+                if label == "Ancient":
+                    labels["ancient"][0] += 1
+                    labels["ancient"][1] += f"{source}:\n{explain}\n\n"
+                elif label == "Modern":
+                    labels["modern"][0] += 1
+                    labels["modern"][1] += f"{source}:\n{explain}\n\n"
+                else:
+                    labels["unknown"] += 1
+
+            if max(labels["modern"][0],labels["ancient"][0]) > 0:
+                if labels["modern"][0] > labels["ancient"][0]:
+                    return "Modern", labels["modern"][1]
+                else:
+                    return "Ancient", labels["ancient"][1]
+            else:
+              return "Unknown", "No strong keywords detected"
+        else:
+            print("No DOI or PubMed ID available for inference.")
+            return "", ""
+
+    except Exception as e:
+        print("Error:", e)
+        return "", ""
+
+
+def detect_ancient_flag(context_snippet):
+    context = context_snippet.lower()
+
+    ancient_keywords = [
+        "ancient", "archaeological", "prehistoric", "neolithic", "mesolithic", "paleolithic",
+        "bronze age", "iron age", "burial", "tomb", "skeleton", "14c", "radiocarbon", "carbon dating",
+        "postmortem damage", "udg treatment", "adna", "degradation", "site", "excavation",
+        "archaeological context", "temporal transect", "population replacement", "cal bp", "calbp", "carbon dated"
+    ]
+
+    modern_keywords = [
+        "modern", "hospital", "clinical", "consent","blood","buccal","unrelated", "blood sample","buccal sample","informed consent", "donor", "healthy", "patient",
+        "genotyping", "screening", "medical", "cohort", "sequencing facility", "ethics approval",
+        "we analysed", "we analyzed", "dataset includes", "new sequences", "published data",
+        "control cohort", "sink population", "genbank accession", "sequenced", "pipeline", 
+        "bioinformatic analysis", "samples from", "population genetics", "genome-wide data", "imr collection"
+    ]
+
+    ancient_hits = [k for k in ancient_keywords if k in context]
+    modern_hits = [k for k in modern_keywords if k in context]
+
+    if ancient_hits and not modern_hits:
+        return "Ancient", f"Flagged as ancient due to keywords: {', '.join(ancient_hits)}"
+    elif modern_hits and not ancient_hits:
+        return "Modern", f"Flagged as modern due to keywords: {', '.join(modern_hits)}"
+    elif ancient_hits and modern_hits:
+        if len(ancient_hits) >= len(modern_hits):
+            return "Ancient", f"Mixed context, leaning ancient due to: {', '.join(ancient_hits)}"
+        else:
+            return "Modern", f"Mixed context, leaning modern due to: {', '.join(modern_hits)}"
+    
+    # Fallback to QA
+    answer = infer_fromQAModel(context, question="Are the mtDNA samples ancient or modern? Explain why.")
+    if answer.startswith("Error"):
+        return "Unknown", answer
+    if "ancient" in answer.lower():
+        return "Ancient", f"Leaning ancient based on QA: {answer}"
+    elif "modern" in answer.lower():
+        return "Modern", f"Leaning modern based on QA: {answer}"
+    else:
+        return "Unknown", f"No strong keywords or QA clues. QA said: {answer}"
+
+# STEP 5: Main pipeline: accession -> 1. get pubmed id and isolate -> 2. get doi -> 3. get text -> 4. prediction -> 5. output: inferred location + explanation + confidence score
+def classify_sample_location(accession):
+  outputs = {}
+  keyword, context, location, qa_result, haplo_result = "", "", "", "", ""
+  # Step 1: get pubmed id and isolate
+  pubmedID, isolate = get_info_from_accession(accession)
+  '''if not pubmedID:
+    return {"error": f"Could not retrieve PubMed ID for accession {accession}"}'''
+  if not isolate:
+    isolate = "UNKNOWN_ISOLATE"
+  # Step 2: get doi
+  doi = get_doi_from_pubmed_id(pubmedID)
+  '''if not doi:
+    return {"error": "DOI not found for this accession. Cannot fetch paper or context."}'''
+  # Step 3: get text
+  '''textsToExtract = { "doiLink":"paperText"
+                        "file1.pdf":"text1",
+                        "file2.doc":"text2",
+                        "file3.xlsx":excelText3'''
+  if doi and pubmedID:                      
+    textsToExtract = get_paper_text(doi,pubmedID)
+  else: textsToExtract = {}  
+  '''if not textsToExtract:
+    return {"error": f"No texts extracted for DOI {doi}"}'''
+  if isolate not in [None, "UNKNOWN_ISOLATE"]:
+    label, explain = flag_ancient_modern(accession,textsToExtract,isolate)
+  else: 
+    label, explain = flag_ancient_modern(accession,textsToExtract)  
+  # Step 4: prediction
+  outputs[accession] = {}
+  outputs[isolate] = {}
+  # 4.0 Infer from NCBI
+  location, outputNCBI = infer_location_fromNCBI(accession)
+  NCBI_result = {
+      "source": "NCBI",
+      "sample_id": accession,
+      "predicted_location": location,
+      "context_snippet": outputNCBI}
+  outputs[accession]["NCBI"]= {"NCBI": NCBI_result}
+  if textsToExtract:
+    long_text = ""
+    for key in textsToExtract:
+      text = textsToExtract[key]
+      # try accession number first
+      outputs[accession][key] = {}
+      keyword = accession
+      context = extract_context(text, keyword, window=500)
+      # 4.1: Using a HuggingFace model (question-answering)
+      location = infer_fromQAModel(context, question=f"Where is the mtDNA sample {keyword} from?")
+      qa_result = {
+          "source": key,
+          "sample_id": keyword,
+          "predicted_location": location,
+          "context_snippet": context
+      }
+      outputs[keyword][key]["QAModel"] = qa_result
+      # 4.2: Infer from haplogroup
+      haplo_result = classify_mtDNA_sample_from_haplo(context)
+      outputs[keyword][key]["haplogroup"] = haplo_result
+      # try isolate
+      keyword = isolate
+      outputs[isolate][key] = {}
+      context = extract_context(text, keyword, window=500)
+      # 4.1.1: Using a HuggingFace model (question-answering)
+      location = infer_fromQAModel(context, question=f"Where is the mtDNA sample {keyword} from?")
+      qa_result = {
+          "source": key,
+          "sample_id": keyword,
+          "predicted_location": location,
+          "context_snippet": context
+      }
+      outputs[keyword][key]["QAModel"] = qa_result
+      # 4.2.1: Infer from haplogroup
+      haplo_result = classify_mtDNA_sample_from_haplo(context)
+      outputs[keyword][key]["haplogroup"] = haplo_result
+      # add long text
+      long_text += text + ". \n"
+    # 4.3: UpgradeClassify
+    # try sample_id as accession number
+    sample_id = accession
+    if sample_id:
+      filtered_context = filter_context_for_sample(sample_id.upper(), long_text, window_size=1)
+      locations = infer_location_for_sample(sample_id.upper(), filtered_context)
+      if locations!="No clear location found in top matches":
+        outputs[sample_id]["upgradeClassifier"] = {}
+        outputs[sample_id]["upgradeClassifier"]["upgradeClassifier"] = {
+          "source": "From these sources combined: "+ ", ".join(list(textsToExtract.keys())),
+          "sample_id": sample_id,
+          "predicted_location": ", ".join(locations),
+          "context_snippep": "First 1000 words: \n"+ filtered_context[:1000]
+        }
+    # try sample_id as isolate name
+    sample_id = isolate
+    if sample_id:
+      filtered_context = filter_context_for_sample(sample_id.upper(), long_text, window_size=1)
+      locations = infer_location_for_sample(sample_id.upper(), filtered_context)
+      if locations!="No clear location found in top matches":
+        outputs[sample_id]["upgradeClassifier"] = {}
+        outputs[sample_id]["upgradeClassifier"]["upgradeClassifier"] = {
+          "source": "From these sources combined: "+ ", ".join(list(textsToExtract.keys())),
+          "sample_id": sample_id,
+          "predicted_location": ", ".join(locations),
+          "context_snippep": "First 1000 words: \n"+ filtered_context[:1000]
+        }
+  return outputs, label, explain

pipeline.py

@@ -0,0 +1,347 @@
+# test1: MJ17 direct
+# test2: "A1YU101" thailand cross-ref
+# test3: "EBK109" thailand cross-ref
+# test4: "OQ731952"/"BST115" for search query title: "South Asian maternal and paternal lineages in southern Thailand and"
+from iterate3 import data_preprocess, model
+import mtdna_classifier
+import app
+import pandas as pd
+from pathlib import Path
+import subprocess
+from NER.html import extractHTML
+import os
+import google.generativeai as genai
+import re
+import standardize_location
+# Helper functions in for this pipeline
+# Track time
+import time
+import multiprocessing
+
+def run_with_timeout(func, args=(), kwargs={}, timeout=20):
+    """
+    Runs `func` with timeout in seconds. Kills if it exceeds.
+    Returns: (success, result or None)
+    """
+    def wrapper(q, *args, **kwargs):
+        try:
+            q.put(func(*args, **kwargs))
+        except Exception as e:
+            q.put(e)
+
+    q = multiprocessing.Queue()
+    p = multiprocessing.Process(target=wrapper, args=(q, *args), kwargs=kwargs)
+    p.start()
+    p.join(timeout)
+
+    if p.is_alive():
+        p.terminate()
+        p.join()
+        print(f"⏱️ Timeout exceeded ({timeout} sec) — function killed.")
+        return False, None
+    else:
+        result = q.get()
+        if isinstance(result, Exception):
+            raise result
+        return True, result
+
+def time_it(func, *args, **kwargs):
+    """
+    Measure how long a function takes to run and return its result + time.
+    """
+    start = time.time()
+    result = func(*args, **kwargs)
+    end = time.time()
+    elapsed = end - start
+    print(f"⏱️ '{func.__name__}' took {elapsed:.3f} seconds")
+    return result, elapsed
+# --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) ---    
+def track_gemini_cost():
+  # Prices are per 1,000 tokens
+  PRICE_PER_1K_INPUT_LLM = 0.000075  # $0.075 per 1M tokens
+  PRICE_PER_1K_OUTPUT_LLM = 0.0003   # $0.30 per 1M tokens
+  PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
+  return True
+
+def unique_preserve_order(seq):
+    seen = set()
+    return [x for x in seq if not (x in seen or seen.add(x))]
+# Main execution
+def pipeline_with_gemini(accessions):
+  # output: country, sample_type, ethnic, location, money_cost, time_cost, explain
+  # there can be one accession number in the accessions
+  # Prices are per 1,000 tokens
+  PRICE_PER_1K_INPUT_LLM = 0.000075  # $0.075 per 1M tokens
+  PRICE_PER_1K_OUTPUT_LLM = 0.0003   # $0.30 per 1M tokens
+  PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
+  if not accessions:
+    print("no input")
+    return None
+  else:  
+    accs_output = {}
+    os.environ["GOOGLE_API_KEY"] = "AIzaSyDi0CNKBgEtnr6YuPaY6YNEuC5wT0cdKhk"
+    genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+    for acc in accessions:
+      start = time.time()
+      total_cost_title = 0
+      jsonSM, links, article_text = {},[], ""
+      acc_score = { "isolate": "",
+                    "country":{},
+                   "sample_type":{},
+                   #"specific_location":{},
+                   #"ethnicity":{},
+                   "query_cost":total_cost_title,
+                   "time_cost":None,
+                   "source":links}
+      meta = mtdna_classifier.fetch_ncbi_metadata(acc)
+      country, spe_loc, ethnic, sample_type, col_date, iso, title, doi, pudID, features = meta["country"], meta["specific_location"], meta["ethnicity"], meta["sample_type"], meta["collection_date"], meta["isolate"], meta["title"], meta["doi"], meta["pubmed_id"], meta["all_features"]
+      acc_score["isolate"] = iso
+      # set up step: create the folder to save document
+      chunk, all_output = "",""
+      if pudID: 
+        id = pudID
+        saveTitle = title
+      else: 
+        saveTitle = title + "_" + col_date
+        id = "DirectSubmission"
+      folder_path = Path("/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/data/"+str(id))
+      if not folder_path.exists():
+          cmd = f'mkdir /content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/data/{id}'
+          result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+          print("data/"+str(id) +" created.")
+      else:
+          print("data/"+str(id) +" already exists.")
+      saveLinkFolder = "/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/data/"+str(id)
+      # first way: ncbi method
+      if country.lower() != "unknown":
+        stand_country = standardize_location.smart_country_lookup(country.lower())
+        if stand_country.lower() != "not found":
+          acc_score["country"][stand_country.lower()] = ["ncbi"]
+        else: acc_score["country"][country.lower()] = ["ncbi"]   
+      # if spe_loc.lower() != "unknown":
+      #   acc_score["specific_location"][spe_loc.lower()] = ["ncbi"]
+      # if ethnic.lower() != "unknown":
+      #   acc_score["ethnicity"][ethnic.lower()] = ["ncbi"]
+      if sample_type.lower() != "unknown":
+        acc_score["sample_type"][sample_type.lower()] = ["ncbi"]
+      # second way: LLM model
+      # Preprocess the input token
+      accession, isolate = None, None
+      if acc != "unknown":  accession = acc
+      if iso != "unknown":  isolate = iso
+      # check doi first
+      if doi != "unknown":
+        link = 'https://doi.org/' + doi
+        # get the file to create listOfFile for each id
+        html = extractHTML.HTML("",link)
+        jsonSM = html.getSupMaterial()
+        article_text = html.getListSection()
+        if article_text:
+          if "Just a moment...Enable JavaScript and cookies to continue".lower() not in article_text.lower() or "403 Forbidden Request".lower() not in article_text.lower():
+            links.append(link)
+        if jsonSM:
+          links += sum((jsonSM[key] for key in jsonSM),[])
+      # no doi then google custom search api
+      if len(article_text) == 0 or "Just a moment...Enable JavaScript and cookies to continue".lower() in article_text.lower() or "403 Forbidden Request".lower() in article_text.lower():
+        # might find the article
+        tem_links = mtdna_classifier.search_google_custom(title, 2)
+        # get supplementary of that article
+        for link in tem_links:
+          html = extractHTML.HTML("",link)
+          jsonSM = html.getSupMaterial()
+          article_text_tem = html.getListSection()
+          if article_text_tem:
+            if "Just a moment...Enable JavaScript and cookies to continue".lower() not in article_text_tem.lower() or "403 Forbidden Request".lower() not in article_text_tem.lower():
+              links.append(link)
+          if jsonSM:
+            links += sum((jsonSM[key] for key in jsonSM),[])
+      print(links)    
+      links = unique_preserve_order(links)
+      acc_score["source"] = links
+      chunk_path = "/"+saveTitle+"_merged_document.docx"
+      all_path = "/"+saveTitle+"_all_merged_document.docx"
+      # if chunk and all output not exist yet
+      file_chunk_path = saveLinkFolder + chunk_path
+      file_all_path = saveLinkFolder + all_path
+      if os.path.exists(file_chunk_path):
+        print("File chunk exists!")
+        if not chunk:
+          text, table, document_title = model.read_docx_text(file_chunk_path)
+          chunk = data_preprocess.normalize_for_overlap(text) + "\n" + data_preprocess.normalize_for_overlap(". ".join(table))
+      if os.path.exists(file_all_path):
+        print("File all output exists!")
+        if not all_output:
+          text_all, table_all, document_title_all = model.read_docx_text(file_all_path)
+          all_output = data_preprocess.normalize_for_overlap(text_all) + "\n" + data_preprocess.normalize_for_overlap(". ".join(table_all))
+      if not chunk and not all_output:
+        # else: check if we can reuse these chunk and all output of existed accession to find another
+        if links:
+          for link in links:
+              print(link)
+              # if len(all_output) > 1000*1000:
+              #   all_output = data_preprocess.normalize_for_overlap(all_output)
+              #   print("after normalizing all output: ", len(all_output))
+              if len(data_preprocess.normalize_for_overlap(all_output)) > 600000:
+                print("break here")
+                break
+              if iso != "unknown": query_kw = iso
+              else: query_kw = acc
+              #text_link, tables_link, final_input_link = data_preprocess.preprocess_document(link,saveLinkFolder, isolate=query_kw)
+              success_process, output_process = run_with_timeout(data_preprocess.preprocess_document,args=(link,saveLinkFolder),kwargs={"isolate":query_kw},timeout=180)
+              if success_process:
+                text_link, tables_link, final_input_link = output_process[0], output_process[1], output_process[2]
+                print("yes succeed for process document")
+              else: text_link, tables_link, final_input_link = "", "", ""  
+              context = data_preprocess.extract_context(final_input_link, query_kw)
+              if context !=  "Sample ID not found.":
+                if len(data_preprocess.normalize_for_overlap(chunk)) < 1000*1000:
+                  success_chunk, the_output_chunk = run_with_timeout(data_preprocess.merge_texts_skipping_overlap,args=(chunk, context))
+                  if success_chunk:
+                    chunk = the_output_chunk#data_preprocess.merge_texts_skipping_overlap(all_output, final_input_link)
+                    print("yes succeed for chunk")
+                  else:
+                    chunk += context
+                    print("len context: ", len(context))
+                    print("basic fall back")
+                print("len chunk after: ", len(chunk))
+              if len(final_input_link) > 1000*1000:
+                if context !=  "Sample ID not found.":
+                  final_input_link =  context 
+                else:
+                  final_input_link = data_preprocess.normalize_for_overlap(final_input_link)
+                  if len(final_input_link) > 1000 *1000:
+                    final_input_link = final_input_link[:100000] 
+              if len(data_preprocess.normalize_for_overlap(all_output)) < 1000*1000:
+                success, the_output = run_with_timeout(data_preprocess.merge_texts_skipping_overlap,args=(all_output, final_input_link))
+                if success:
+                  all_output = the_output#data_preprocess.merge_texts_skipping_overlap(all_output, final_input_link)
+                  print("yes succeed")
+                else:
+                  all_output += final_input_link
+                  print("len final input: ", len(final_input_link))
+                  print("basic fall back")
+              print("len all output after: ", len(all_output))
+          #country_pro, chunk, all_output = data_preprocess.process_inputToken(links, saveLinkFolder, accession=accession, isolate=isolate)
+          
+        else:
+          chunk = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
+          all_output = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
+        if not chunk: chunk = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
+        if not all_output:  all_output = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
+        if len(all_output) > 1*1024*1024: 
+          all_output = data_preprocess.normalize_for_overlap(all_output)
+          if len(all_output) > 1*1024*1024:
+            all_output = all_output[:1*1024*1024]
+        print("chunk len: ", len(chunk))
+        print("all output len: ", len(all_output))    
+        data_preprocess.save_text_to_docx(chunk, file_chunk_path)
+        data_preprocess.save_text_to_docx(all_output, file_all_path)
+      # else:
+      #   final_input = ""
+      #   if all_output:
+      #     final_input = all_output
+      #   else:  
+      #     if chunk: final_input = chunk
+      #   #data_preprocess.merge_texts_skipping_overlap(final_input, all_output)
+      #   if final_input:
+      #     keywords = []
+      #     if iso != "unknown":  keywords.append(iso)
+      #     if acc != "unknown":  keywords.append(acc)
+      #     for keyword in keywords:
+      #       chunkBFS = data_preprocess.get_contextual_sentences_BFS(final_input, keyword)
+      #       countryDFS, chunkDFS = data_preprocess.get_contextual_sentences_DFS(final_input, keyword)
+      #       chunk = data_preprocess.merge_texts_skipping_overlap(chunk, chunkDFS)
+      #       chunk = data_preprocess.merge_texts_skipping_overlap(chunk, chunkBFS)
+          
+      # Define paths for cached RAG assets
+      faiss_index_path = saveLinkFolder+"/faiss_index.bin"
+      document_chunks_path = saveLinkFolder+"/document_chunks.json"
+      structured_lookup_path = saveLinkFolder+"/structured_lookup.json"
+
+      master_structured_lookup, faiss_index, document_chunks = model.load_rag_assets(
+          faiss_index_path, document_chunks_path, structured_lookup_path
+      )
+
+      global_llm_model_for_counting_tokens = genai.GenerativeModel('gemini-1.5-flash-latest')
+      if not all_output:
+        if chunk: all_output = chunk
+        else: all_output = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
+      if faiss_index is None:
+          print("\nBuilding RAG assets (structured lookup, FAISS index, chunks)...")
+          total_doc_embedding_tokens = global_llm_model_for_counting_tokens.count_tokens(
+              all_output
+          ).total_tokens
+
+          initial_embedding_cost = (total_doc_embedding_tokens / 1000) * PRICE_PER_1K_EMBEDDING_INPUT
+          total_cost_title += initial_embedding_cost
+          print(f"Initial one-time embedding cost for '{file_all_path}' ({total_doc_embedding_tokens} tokens): ${initial_embedding_cost:.6f}")
+
+
+          master_structured_lookup, faiss_index, document_chunks, plain_text_content = model.build_vector_index_and_data(
+              file_all_path, faiss_index_path, document_chunks_path, structured_lookup_path
+          )
+      else:
+          print("\nRAG assets loaded from file. No re-embedding of entire document will occur.")
+          plain_text_content_all, table_strings_all, document_title_all = model.read_docx_text(file_all_path)
+          master_structured_lookup['document_title'] = master_structured_lookup.get('document_title', document_title_all)
+
+      primary_word = iso
+      alternative_word = acc
+      print(f"\n--- General Query: Primary='{primary_word}' (Alternative='{alternative_word}') ---")
+      if features.lower() not in all_output.lower():  
+        all_output += ". NCBI Features: " + features
+      # country, sample_type, method_used, ethnic, spe_loc, total_query_cost =  model.query_document_info(
+      #     primary_word, alternative_word, meta, master_structured_lookup, faiss_index, document_chunks, 
+      #     model.call_llm_api, chunk=chunk, all_output=all_output)
+      country, sample_type, method_used, country_explanation, sample_type_explanation, total_query_cost =  model.query_document_info(
+          primary_word, alternative_word, meta, master_structured_lookup, faiss_index, document_chunks, 
+          model.call_llm_api, chunk=chunk, all_output=all_output)
+      if len(country) == 0: country = "unknown"
+      if len(sample_type) == 0: sample_type = "unknown"    
+      if country_explanation: country_explanation = "-"+country_explanation        
+      else: country_explanation = ""
+      if sample_type_explanation: sample_type_explanation = "-"+sample_type_explanation
+      else: sample_type_explanation = ""
+      if method_used == "unknown": method_used = ""
+      if country.lower() != "unknown":
+        stand_country = standardize_location.smart_country_lookup(country.lower())
+        if stand_country.lower() != "not found":
+          if stand_country.lower() in acc_score["country"]:
+            if country_explanation:
+              acc_score["country"][stand_country.lower()].append(method_used + country_explanation)
+          else:
+            acc_score["country"][stand_country.lower()] = [method_used + country_explanation]
+        else:
+          if country.lower() in acc_score["country"]:
+            if country_explanation:
+              if len(method_used + country_explanation) > 0:
+                acc_score["country"][country.lower()].append(method_used + country_explanation)
+          else:
+            if len(method_used + country_explanation) > 0:
+              acc_score["country"][country.lower()] = [method_used + country_explanation]
+      # if spe_loc.lower() != "unknown":
+      #   if spe_loc.lower() in acc_score["specific_location"]:
+      #     acc_score["specific_location"][spe_loc.lower()].append(method_used)
+      #   else:
+      #     acc_score["specific_location"][spe_loc.lower()] = [method_used]
+      # if ethnic.lower() != "unknown":
+      #   if ethnic.lower() in acc_score["ethnicity"]:
+      #     acc_score["ethnicity"][ethnic.lower()].append(method_used)
+      #   else:
+      #     acc_score["ethnicity"][ethnic.lower()] = [method_used]
+      if sample_type.lower() != "unknown":
+        if sample_type.lower() in acc_score["sample_type"]:
+          if len(method_used + sample_type_explanation) > 0:
+            acc_score["sample_type"][sample_type.lower()].append(method_used + sample_type_explanation)
+        else:
+          if len(method_used + sample_type_explanation)> 0:
+            acc_score["sample_type"][sample_type.lower()] = [method_used + sample_type_explanation]
+      end = time.time()
+      total_cost_title += total_query_cost
+      acc_score["query_cost"] = total_cost_title
+      elapsed = end - start
+      acc_score["time_cost"] = f"{elapsed:.3f} seconds"
+      accs_output[acc] = acc_score
+      print(accs_output[acc])
+      
+    return accs_output

requirements.txt

@@ -1,7 +1,7 @@
 biopython==1.85
 bs4==0.0.2
 gensim==4.3.3
-gradio==5.29.0
+gradio
 gspread==6.2.0
 gspread-dataframe==4.0.0
 huggingface-hub==0.30.2
@@ -23,9 +23,22 @@ Spire.Xls==14.12.0
 statsmodels==0.14.4
 tabula-py==2.10.0
 thefuzz==0.22.1
-torch @ https://download.pytorch.org/whl/cu124/torch-2.6.0%2Bcu124-cp310-cp310-linux_x86_64.whl
+torch
 transformers==4.51.3
 wordsegment==1.3.1
 xlrd==2.0.1
 sentence-transformers
-lxml
+lxml
+streamlit
+requests
+google-generativeai
+PyPDF2
+beautifulsoup4
+# For Claude
+anthropic
+faiss-cpu
+python-docx
+pycountry
+# For Deepseek (If direct DeepseekLLM client library is available, use it.
+# Otherwise, 'requests' covers it for simple API calls, but a dedicated client is better for full features)
+# deepseek-llm # Uncomment this if Deepseek provides a dedicated pip package for their LLM