Spaces:
Running
Running
Update pipeline.py
Browse files- pipeline.py +345 -346
pipeline.py
CHANGED
@@ -1,347 +1,346 @@
|
|
1 |
-
# test1: MJ17 direct
|
2 |
-
# test2: "A1YU101" thailand cross-ref
|
3 |
-
# test3: "EBK109" thailand cross-ref
|
4 |
-
# test4: "OQ731952"/"BST115" for search query title: "South Asian maternal and paternal lineages in southern Thailand and"
|
5 |
-
from iterate3 import data_preprocess, model
|
6 |
-
import mtdna_classifier
|
7 |
-
import app
|
8 |
-
import pandas as pd
|
9 |
-
from pathlib import Path
|
10 |
-
import subprocess
|
11 |
-
from NER.html import extractHTML
|
12 |
-
import os
|
13 |
-
import google.generativeai as genai
|
14 |
-
import re
|
15 |
-
import standardize_location
|
16 |
-
# Helper functions in for this pipeline
|
17 |
-
# Track time
|
18 |
-
import time
|
19 |
-
import multiprocessing
|
20 |
-
|
21 |
-
def run_with_timeout(func, args=(), kwargs={}, timeout=20):
|
22 |
-
"""
|
23 |
-
Runs `func` with timeout in seconds. Kills if it exceeds.
|
24 |
-
Returns: (success, result or None)
|
25 |
-
"""
|
26 |
-
def wrapper(q, *args, **kwargs):
|
27 |
-
try:
|
28 |
-
q.put(func(*args, **kwargs))
|
29 |
-
except Exception as e:
|
30 |
-
q.put(e)
|
31 |
-
|
32 |
-
q = multiprocessing.Queue()
|
33 |
-
p = multiprocessing.Process(target=wrapper, args=(q, *args), kwargs=kwargs)
|
34 |
-
p.start()
|
35 |
-
p.join(timeout)
|
36 |
-
|
37 |
-
if p.is_alive():
|
38 |
-
p.terminate()
|
39 |
-
p.join()
|
40 |
-
print(f"⏱️ Timeout exceeded ({timeout} sec) — function killed.")
|
41 |
-
return False, None
|
42 |
-
else:
|
43 |
-
result = q.get()
|
44 |
-
if isinstance(result, Exception):
|
45 |
-
raise result
|
46 |
-
return True, result
|
47 |
-
|
48 |
-
def time_it(func, *args, **kwargs):
|
49 |
-
"""
|
50 |
-
Measure how long a function takes to run and return its result + time.
|
51 |
-
"""
|
52 |
-
start = time.time()
|
53 |
-
result = func(*args, **kwargs)
|
54 |
-
end = time.time()
|
55 |
-
elapsed = end - start
|
56 |
-
print(f"⏱️ '{func.__name__}' took {elapsed:.3f} seconds")
|
57 |
-
return result, elapsed
|
58 |
-
# --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) ---
|
59 |
-
def track_gemini_cost():
|
60 |
-
# Prices are per 1,000 tokens
|
61 |
-
PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens
|
62 |
-
PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens
|
63 |
-
PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
|
64 |
-
return True
|
65 |
-
|
66 |
-
def unique_preserve_order(seq):
|
67 |
-
seen = set()
|
68 |
-
return [x for x in seq if not (x in seen or seen.add(x))]
|
69 |
-
# Main execution
|
70 |
-
def pipeline_with_gemini(accessions):
|
71 |
-
# output: country, sample_type, ethnic, location, money_cost, time_cost, explain
|
72 |
-
# there can be one accession number in the accessions
|
73 |
-
# Prices are per 1,000 tokens
|
74 |
-
PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens
|
75 |
-
PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens
|
76 |
-
PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
|
77 |
-
if not accessions:
|
78 |
-
print("no input")
|
79 |
-
return None
|
80 |
-
else:
|
81 |
-
accs_output = {}
|
82 |
-
os.
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
"
|
91 |
-
#"
|
92 |
-
|
93 |
-
"
|
94 |
-
"
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
cmd =
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
stand_country
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
#
|
122 |
-
#
|
123 |
-
#
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
#
|
128 |
-
|
129 |
-
accession
|
130 |
-
if
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
article_text
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
article_text_tem
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
text
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
text_all
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
#
|
182 |
-
#
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
189 |
-
|
190 |
-
success_process
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
context
|
196 |
-
|
197 |
-
|
198 |
-
success_chunk
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
print("
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
success
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
print("
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
229 |
-
if not
|
230 |
-
if
|
231 |
-
|
232 |
-
|
233 |
-
|
234 |
-
|
235 |
-
print("
|
236 |
-
|
237 |
-
data_preprocess.save_text_to_docx(
|
238 |
-
|
239 |
-
#
|
240 |
-
#
|
241 |
-
#
|
242 |
-
#
|
243 |
-
#
|
244 |
-
#
|
245 |
-
#
|
246 |
-
#
|
247 |
-
#
|
248 |
-
# if
|
249 |
-
#
|
250 |
-
#
|
251 |
-
#
|
252 |
-
#
|
253 |
-
# chunk = data_preprocess.merge_texts_skipping_overlap(chunk,
|
254 |
-
|
255 |
-
|
256 |
-
|
257 |
-
|
258 |
-
|
259 |
-
|
260 |
-
|
261 |
-
|
262 |
-
|
263 |
-
|
264 |
-
|
265 |
-
|
266 |
-
|
267 |
-
|
268 |
-
|
269 |
-
|
270 |
-
|
271 |
-
|
272 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
|
281 |
-
|
282 |
-
|
283 |
-
|
284 |
-
|
285 |
-
|
286 |
-
|
287 |
-
|
288 |
-
|
289 |
-
|
290 |
-
|
291 |
-
|
292 |
-
|
293 |
-
#
|
294 |
-
#
|
295 |
-
|
296 |
-
|
297 |
-
|
298 |
-
|
299 |
-
if len(
|
300 |
-
if
|
301 |
-
|
302 |
-
|
303 |
-
|
304 |
-
|
305 |
-
if
|
306 |
-
|
307 |
-
stand_country
|
308 |
-
|
309 |
-
|
310 |
-
|
311 |
-
|
312 |
-
|
313 |
-
|
314 |
-
|
315 |
-
|
316 |
-
|
317 |
-
|
318 |
-
|
319 |
-
|
320 |
-
|
321 |
-
|
322 |
-
#
|
323 |
-
#
|
324 |
-
#
|
325 |
-
#
|
326 |
-
#
|
327 |
-
#
|
328 |
-
#
|
329 |
-
#
|
330 |
-
#
|
331 |
-
|
332 |
-
|
333 |
-
|
334 |
-
|
335 |
-
|
336 |
-
|
337 |
-
|
338 |
-
|
339 |
-
|
340 |
-
|
341 |
-
|
342 |
-
|
343 |
-
|
344 |
-
accs_output[acc]
|
345 |
-
|
346 |
-
|
347 |
return accs_output
|
|
|
1 |
+
# test1: MJ17 direct
|
2 |
+
# test2: "A1YU101" thailand cross-ref
|
3 |
+
# test3: "EBK109" thailand cross-ref
|
4 |
+
# test4: "OQ731952"/"BST115" for search query title: "South Asian maternal and paternal lineages in southern Thailand and"
|
5 |
+
from iterate3 import data_preprocess, model
|
6 |
+
import mtdna_classifier
|
7 |
+
import app
|
8 |
+
import pandas as pd
|
9 |
+
from pathlib import Path
|
10 |
+
import subprocess
|
11 |
+
from NER.html import extractHTML
|
12 |
+
import os
|
13 |
+
import google.generativeai as genai
|
14 |
+
import re
|
15 |
+
import standardize_location
|
16 |
+
# Helper functions in for this pipeline
|
17 |
+
# Track time
|
18 |
+
import time
|
19 |
+
import multiprocessing
|
20 |
+
|
21 |
+
def run_with_timeout(func, args=(), kwargs={}, timeout=20):
|
22 |
+
"""
|
23 |
+
Runs `func` with timeout in seconds. Kills if it exceeds.
|
24 |
+
Returns: (success, result or None)
|
25 |
+
"""
|
26 |
+
def wrapper(q, *args, **kwargs):
|
27 |
+
try:
|
28 |
+
q.put(func(*args, **kwargs))
|
29 |
+
except Exception as e:
|
30 |
+
q.put(e)
|
31 |
+
|
32 |
+
q = multiprocessing.Queue()
|
33 |
+
p = multiprocessing.Process(target=wrapper, args=(q, *args), kwargs=kwargs)
|
34 |
+
p.start()
|
35 |
+
p.join(timeout)
|
36 |
+
|
37 |
+
if p.is_alive():
|
38 |
+
p.terminate()
|
39 |
+
p.join()
|
40 |
+
print(f"⏱️ Timeout exceeded ({timeout} sec) — function killed.")
|
41 |
+
return False, None
|
42 |
+
else:
|
43 |
+
result = q.get()
|
44 |
+
if isinstance(result, Exception):
|
45 |
+
raise result
|
46 |
+
return True, result
|
47 |
+
|
48 |
+
def time_it(func, *args, **kwargs):
|
49 |
+
"""
|
50 |
+
Measure how long a function takes to run and return its result + time.
|
51 |
+
"""
|
52 |
+
start = time.time()
|
53 |
+
result = func(*args, **kwargs)
|
54 |
+
end = time.time()
|
55 |
+
elapsed = end - start
|
56 |
+
print(f"⏱️ '{func.__name__}' took {elapsed:.3f} seconds")
|
57 |
+
return result, elapsed
|
58 |
+
# --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) ---
|
59 |
+
def track_gemini_cost():
|
60 |
+
# Prices are per 1,000 tokens
|
61 |
+
PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens
|
62 |
+
PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens
|
63 |
+
PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
|
64 |
+
return True
|
65 |
+
|
66 |
+
def unique_preserve_order(seq):
|
67 |
+
seen = set()
|
68 |
+
return [x for x in seq if not (x in seen or seen.add(x))]
|
69 |
+
# Main execution
|
70 |
+
def pipeline_with_gemini(accessions):
|
71 |
+
# output: country, sample_type, ethnic, location, money_cost, time_cost, explain
|
72 |
+
# there can be one accession number in the accessions
|
73 |
+
# Prices are per 1,000 tokens
|
74 |
+
PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens
|
75 |
+
PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens
|
76 |
+
PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
|
77 |
+
if not accessions:
|
78 |
+
print("no input")
|
79 |
+
return None
|
80 |
+
else:
|
81 |
+
accs_output = {}
|
82 |
+
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
|
83 |
+
for acc in accessions:
|
84 |
+
start = time.time()
|
85 |
+
total_cost_title = 0
|
86 |
+
jsonSM, links, article_text = {},[], ""
|
87 |
+
acc_score = { "isolate": "",
|
88 |
+
"country":{},
|
89 |
+
"sample_type":{},
|
90 |
+
#"specific_location":{},
|
91 |
+
#"ethnicity":{},
|
92 |
+
"query_cost":total_cost_title,
|
93 |
+
"time_cost":None,
|
94 |
+
"source":links}
|
95 |
+
meta = mtdna_classifier.fetch_ncbi_metadata(acc)
|
96 |
+
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"]
|
97 |
+
acc_score["isolate"] = iso
|
98 |
+
# set up step: create the folder to save document
|
99 |
+
chunk, all_output = "",""
|
100 |
+
if pudID:
|
101 |
+
id = pudID
|
102 |
+
saveTitle = title
|
103 |
+
else:
|
104 |
+
saveTitle = title + "_" + col_date
|
105 |
+
id = "DirectSubmission"
|
106 |
+
folder_path = Path("/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/data/"+str(id))
|
107 |
+
if not folder_path.exists():
|
108 |
+
cmd = f'mkdir /content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/data/{id}'
|
109 |
+
result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
110 |
+
print("data/"+str(id) +" created.")
|
111 |
+
else:
|
112 |
+
print("data/"+str(id) +" already exists.")
|
113 |
+
saveLinkFolder = "/content/drive/MyDrive/CollectData/MVP/mtDNA-Location-Classifier/data/"+str(id)
|
114 |
+
# first way: ncbi method
|
115 |
+
if country.lower() != "unknown":
|
116 |
+
stand_country = standardize_location.smart_country_lookup(country.lower())
|
117 |
+
if stand_country.lower() != "not found":
|
118 |
+
acc_score["country"][stand_country.lower()] = ["ncbi"]
|
119 |
+
else: acc_score["country"][country.lower()] = ["ncbi"]
|
120 |
+
# if spe_loc.lower() != "unknown":
|
121 |
+
# acc_score["specific_location"][spe_loc.lower()] = ["ncbi"]
|
122 |
+
# if ethnic.lower() != "unknown":
|
123 |
+
# acc_score["ethnicity"][ethnic.lower()] = ["ncbi"]
|
124 |
+
if sample_type.lower() != "unknown":
|
125 |
+
acc_score["sample_type"][sample_type.lower()] = ["ncbi"]
|
126 |
+
# second way: LLM model
|
127 |
+
# Preprocess the input token
|
128 |
+
accession, isolate = None, None
|
129 |
+
if acc != "unknown": accession = acc
|
130 |
+
if iso != "unknown": isolate = iso
|
131 |
+
# check doi first
|
132 |
+
if doi != "unknown":
|
133 |
+
link = 'https://doi.org/' + doi
|
134 |
+
# get the file to create listOfFile for each id
|
135 |
+
html = extractHTML.HTML("",link)
|
136 |
+
jsonSM = html.getSupMaterial()
|
137 |
+
article_text = html.getListSection()
|
138 |
+
if article_text:
|
139 |
+
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():
|
140 |
+
links.append(link)
|
141 |
+
if jsonSM:
|
142 |
+
links += sum((jsonSM[key] for key in jsonSM),[])
|
143 |
+
# no doi then google custom search api
|
144 |
+
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():
|
145 |
+
# might find the article
|
146 |
+
tem_links = mtdna_classifier.search_google_custom(title, 2)
|
147 |
+
# get supplementary of that article
|
148 |
+
for link in tem_links:
|
149 |
+
html = extractHTML.HTML("",link)
|
150 |
+
jsonSM = html.getSupMaterial()
|
151 |
+
article_text_tem = html.getListSection()
|
152 |
+
if article_text_tem:
|
153 |
+
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():
|
154 |
+
links.append(link)
|
155 |
+
if jsonSM:
|
156 |
+
links += sum((jsonSM[key] for key in jsonSM),[])
|
157 |
+
print(links)
|
158 |
+
links = unique_preserve_order(links)
|
159 |
+
acc_score["source"] = links
|
160 |
+
chunk_path = "/"+saveTitle+"_merged_document.docx"
|
161 |
+
all_path = "/"+saveTitle+"_all_merged_document.docx"
|
162 |
+
# if chunk and all output not exist yet
|
163 |
+
file_chunk_path = saveLinkFolder + chunk_path
|
164 |
+
file_all_path = saveLinkFolder + all_path
|
165 |
+
if os.path.exists(file_chunk_path):
|
166 |
+
print("File chunk exists!")
|
167 |
+
if not chunk:
|
168 |
+
text, table, document_title = model.read_docx_text(file_chunk_path)
|
169 |
+
chunk = data_preprocess.normalize_for_overlap(text) + "\n" + data_preprocess.normalize_for_overlap(". ".join(table))
|
170 |
+
if os.path.exists(file_all_path):
|
171 |
+
print("File all output exists!")
|
172 |
+
if not all_output:
|
173 |
+
text_all, table_all, document_title_all = model.read_docx_text(file_all_path)
|
174 |
+
all_output = data_preprocess.normalize_for_overlap(text_all) + "\n" + data_preprocess.normalize_for_overlap(". ".join(table_all))
|
175 |
+
if not chunk and not all_output:
|
176 |
+
# else: check if we can reuse these chunk and all output of existed accession to find another
|
177 |
+
if links:
|
178 |
+
for link in links:
|
179 |
+
print(link)
|
180 |
+
# if len(all_output) > 1000*1000:
|
181 |
+
# all_output = data_preprocess.normalize_for_overlap(all_output)
|
182 |
+
# print("after normalizing all output: ", len(all_output))
|
183 |
+
if len(data_preprocess.normalize_for_overlap(all_output)) > 600000:
|
184 |
+
print("break here")
|
185 |
+
break
|
186 |
+
if iso != "unknown": query_kw = iso
|
187 |
+
else: query_kw = acc
|
188 |
+
#text_link, tables_link, final_input_link = data_preprocess.preprocess_document(link,saveLinkFolder, isolate=query_kw)
|
189 |
+
success_process, output_process = run_with_timeout(data_preprocess.preprocess_document,args=(link,saveLinkFolder),kwargs={"isolate":query_kw},timeout=180)
|
190 |
+
if success_process:
|
191 |
+
text_link, tables_link, final_input_link = output_process[0], output_process[1], output_process[2]
|
192 |
+
print("yes succeed for process document")
|
193 |
+
else: text_link, tables_link, final_input_link = "", "", ""
|
194 |
+
context = data_preprocess.extract_context(final_input_link, query_kw)
|
195 |
+
if context != "Sample ID not found.":
|
196 |
+
if len(data_preprocess.normalize_for_overlap(chunk)) < 1000*1000:
|
197 |
+
success_chunk, the_output_chunk = run_with_timeout(data_preprocess.merge_texts_skipping_overlap,args=(chunk, context))
|
198 |
+
if success_chunk:
|
199 |
+
chunk = the_output_chunk#data_preprocess.merge_texts_skipping_overlap(all_output, final_input_link)
|
200 |
+
print("yes succeed for chunk")
|
201 |
+
else:
|
202 |
+
chunk += context
|
203 |
+
print("len context: ", len(context))
|
204 |
+
print("basic fall back")
|
205 |
+
print("len chunk after: ", len(chunk))
|
206 |
+
if len(final_input_link) > 1000*1000:
|
207 |
+
if context != "Sample ID not found.":
|
208 |
+
final_input_link = context
|
209 |
+
else:
|
210 |
+
final_input_link = data_preprocess.normalize_for_overlap(final_input_link)
|
211 |
+
if len(final_input_link) > 1000 *1000:
|
212 |
+
final_input_link = final_input_link[:100000]
|
213 |
+
if len(data_preprocess.normalize_for_overlap(all_output)) < 1000*1000:
|
214 |
+
success, the_output = run_with_timeout(data_preprocess.merge_texts_skipping_overlap,args=(all_output, final_input_link))
|
215 |
+
if success:
|
216 |
+
all_output = the_output#data_preprocess.merge_texts_skipping_overlap(all_output, final_input_link)
|
217 |
+
print("yes succeed")
|
218 |
+
else:
|
219 |
+
all_output += final_input_link
|
220 |
+
print("len final input: ", len(final_input_link))
|
221 |
+
print("basic fall back")
|
222 |
+
print("len all output after: ", len(all_output))
|
223 |
+
#country_pro, chunk, all_output = data_preprocess.process_inputToken(links, saveLinkFolder, accession=accession, isolate=isolate)
|
224 |
+
|
225 |
+
else:
|
226 |
+
chunk = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
|
227 |
+
all_output = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
|
228 |
+
if not chunk: chunk = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
|
229 |
+
if not all_output: all_output = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
|
230 |
+
if len(all_output) > 1*1024*1024:
|
231 |
+
all_output = data_preprocess.normalize_for_overlap(all_output)
|
232 |
+
if len(all_output) > 1*1024*1024:
|
233 |
+
all_output = all_output[:1*1024*1024]
|
234 |
+
print("chunk len: ", len(chunk))
|
235 |
+
print("all output len: ", len(all_output))
|
236 |
+
data_preprocess.save_text_to_docx(chunk, file_chunk_path)
|
237 |
+
data_preprocess.save_text_to_docx(all_output, file_all_path)
|
238 |
+
# else:
|
239 |
+
# final_input = ""
|
240 |
+
# if all_output:
|
241 |
+
# final_input = all_output
|
242 |
+
# else:
|
243 |
+
# if chunk: final_input = chunk
|
244 |
+
# #data_preprocess.merge_texts_skipping_overlap(final_input, all_output)
|
245 |
+
# if final_input:
|
246 |
+
# keywords = []
|
247 |
+
# if iso != "unknown": keywords.append(iso)
|
248 |
+
# if acc != "unknown": keywords.append(acc)
|
249 |
+
# for keyword in keywords:
|
250 |
+
# chunkBFS = data_preprocess.get_contextual_sentences_BFS(final_input, keyword)
|
251 |
+
# countryDFS, chunkDFS = data_preprocess.get_contextual_sentences_DFS(final_input, keyword)
|
252 |
+
# chunk = data_preprocess.merge_texts_skipping_overlap(chunk, chunkDFS)
|
253 |
+
# chunk = data_preprocess.merge_texts_skipping_overlap(chunk, chunkBFS)
|
254 |
+
|
255 |
+
# Define paths for cached RAG assets
|
256 |
+
faiss_index_path = saveLinkFolder+"/faiss_index.bin"
|
257 |
+
document_chunks_path = saveLinkFolder+"/document_chunks.json"
|
258 |
+
structured_lookup_path = saveLinkFolder+"/structured_lookup.json"
|
259 |
+
|
260 |
+
master_structured_lookup, faiss_index, document_chunks = model.load_rag_assets(
|
261 |
+
faiss_index_path, document_chunks_path, structured_lookup_path
|
262 |
+
)
|
263 |
+
|
264 |
+
global_llm_model_for_counting_tokens = genai.GenerativeModel('gemini-1.5-flash-latest')
|
265 |
+
if not all_output:
|
266 |
+
if chunk: all_output = chunk
|
267 |
+
else: all_output = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + features
|
268 |
+
if faiss_index is None:
|
269 |
+
print("\nBuilding RAG assets (structured lookup, FAISS index, chunks)...")
|
270 |
+
total_doc_embedding_tokens = global_llm_model_for_counting_tokens.count_tokens(
|
271 |
+
all_output
|
272 |
+
).total_tokens
|
273 |
+
|
274 |
+
initial_embedding_cost = (total_doc_embedding_tokens / 1000) * PRICE_PER_1K_EMBEDDING_INPUT
|
275 |
+
total_cost_title += initial_embedding_cost
|
276 |
+
print(f"Initial one-time embedding cost for '{file_all_path}' ({total_doc_embedding_tokens} tokens): ${initial_embedding_cost:.6f}")
|
277 |
+
|
278 |
+
|
279 |
+
master_structured_lookup, faiss_index, document_chunks, plain_text_content = model.build_vector_index_and_data(
|
280 |
+
file_all_path, faiss_index_path, document_chunks_path, structured_lookup_path
|
281 |
+
)
|
282 |
+
else:
|
283 |
+
print("\nRAG assets loaded from file. No re-embedding of entire document will occur.")
|
284 |
+
plain_text_content_all, table_strings_all, document_title_all = model.read_docx_text(file_all_path)
|
285 |
+
master_structured_lookup['document_title'] = master_structured_lookup.get('document_title', document_title_all)
|
286 |
+
|
287 |
+
primary_word = iso
|
288 |
+
alternative_word = acc
|
289 |
+
print(f"\n--- General Query: Primary='{primary_word}' (Alternative='{alternative_word}') ---")
|
290 |
+
if features.lower() not in all_output.lower():
|
291 |
+
all_output += ". NCBI Features: " + features
|
292 |
+
# country, sample_type, method_used, ethnic, spe_loc, total_query_cost = model.query_document_info(
|
293 |
+
# primary_word, alternative_word, meta, master_structured_lookup, faiss_index, document_chunks,
|
294 |
+
# model.call_llm_api, chunk=chunk, all_output=all_output)
|
295 |
+
country, sample_type, method_used, country_explanation, sample_type_explanation, total_query_cost = model.query_document_info(
|
296 |
+
primary_word, alternative_word, meta, master_structured_lookup, faiss_index, document_chunks,
|
297 |
+
model.call_llm_api, chunk=chunk, all_output=all_output)
|
298 |
+
if len(country) == 0: country = "unknown"
|
299 |
+
if len(sample_type) == 0: sample_type = "unknown"
|
300 |
+
if country_explanation: country_explanation = "-"+country_explanation
|
301 |
+
else: country_explanation = ""
|
302 |
+
if sample_type_explanation: sample_type_explanation = "-"+sample_type_explanation
|
303 |
+
else: sample_type_explanation = ""
|
304 |
+
if method_used == "unknown": method_used = ""
|
305 |
+
if country.lower() != "unknown":
|
306 |
+
stand_country = standardize_location.smart_country_lookup(country.lower())
|
307 |
+
if stand_country.lower() != "not found":
|
308 |
+
if stand_country.lower() in acc_score["country"]:
|
309 |
+
if country_explanation:
|
310 |
+
acc_score["country"][stand_country.lower()].append(method_used + country_explanation)
|
311 |
+
else:
|
312 |
+
acc_score["country"][stand_country.lower()] = [method_used + country_explanation]
|
313 |
+
else:
|
314 |
+
if country.lower() in acc_score["country"]:
|
315 |
+
if country_explanation:
|
316 |
+
if len(method_used + country_explanation) > 0:
|
317 |
+
acc_score["country"][country.lower()].append(method_used + country_explanation)
|
318 |
+
else:
|
319 |
+
if len(method_used + country_explanation) > 0:
|
320 |
+
acc_score["country"][country.lower()] = [method_used + country_explanation]
|
321 |
+
# if spe_loc.lower() != "unknown":
|
322 |
+
# if spe_loc.lower() in acc_score["specific_location"]:
|
323 |
+
# acc_score["specific_location"][spe_loc.lower()].append(method_used)
|
324 |
+
# else:
|
325 |
+
# acc_score["specific_location"][spe_loc.lower()] = [method_used]
|
326 |
+
# if ethnic.lower() != "unknown":
|
327 |
+
# if ethnic.lower() in acc_score["ethnicity"]:
|
328 |
+
# acc_score["ethnicity"][ethnic.lower()].append(method_used)
|
329 |
+
# else:
|
330 |
+
# acc_score["ethnicity"][ethnic.lower()] = [method_used]
|
331 |
+
if sample_type.lower() != "unknown":
|
332 |
+
if sample_type.lower() in acc_score["sample_type"]:
|
333 |
+
if len(method_used + sample_type_explanation) > 0:
|
334 |
+
acc_score["sample_type"][sample_type.lower()].append(method_used + sample_type_explanation)
|
335 |
+
else:
|
336 |
+
if len(method_used + sample_type_explanation)> 0:
|
337 |
+
acc_score["sample_type"][sample_type.lower()] = [method_used + sample_type_explanation]
|
338 |
+
end = time.time()
|
339 |
+
total_cost_title += total_query_cost
|
340 |
+
acc_score["query_cost"] = total_cost_title
|
341 |
+
elapsed = end - start
|
342 |
+
acc_score["time_cost"] = f"{elapsed:.3f} seconds"
|
343 |
+
accs_output[acc] = acc_score
|
344 |
+
print(accs_output[acc])
|
345 |
+
|
|
|
346 |
return accs_output
|