Restructure LLM calls

app.py

@@ -3,24 +3,14 @@ import os
 import pandas as pd
 import streamlit as st
 
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from langchain_community.document_loaders.pdf import PyPDFLoader
 from langchain_core.documents.base import Document
 from langchain_text_splitters import TokenTextSplitter
-from process import get_entity, get_entity_one, get_table, validate
+from process import Process
 from tempfile import NamedTemporaryFile
 from stqdm import stqdm
-from threading import Thread
-
-class CustomThread(Thread):
-    def __init__(self, func, chunk):
-        super().__init__()
-        self.func = func
-        self.chunk = chunk
-        self.result = ''
-
-    def run(self):
-        self.result = self.func(self.chunk)
 
 buffer = io.BytesIO()
 
@@ -34,18 +24,26 @@ uploaded_files = st.file_uploader("Upload Paper(s) here :", type="pdf", accept_m
 col1, col2 = st.columns(2)
 
 with col1:
-    chunk_option = st.selectbox(
-        'Token amounts per process:',
-        (24000, 16000, 8000), key='token'
+    models = (
+        'gpt-4-turbo',
+        'gemini-1.5-pro-latest'
+        # 'llama-3-sonar-large-32k-chat',
+        # 'mixtral-8x7b-instruct',
+    )
+    model = st.selectbox(
+        'Model selection:', models, key='model'
     )
-    chunk_overlap = 0
 
 with col2:
-    model = st.selectbox(
-        'Model selection: (UNDER DEVELOPED)',
-        # 128000, 32768, 1048576
-        ('gpt-4-turbo', 'llama-3-sonar-large-32k-chat', 'gemini-1.5-pro-latest'), key='model'
+    tokens = (
+        24000,
+        16000,
+        8000
     )
+    chunk_option = st.selectbox(
+        'Token amounts per process:', tokens, key='token'
+    )
+    chunk_overlap = 0
 
 if uploaded_files:
     journals = []
@@ -58,6 +56,8 @@ if uploaded_files:
             for uploaded_file in stqdm(uploaded_files):
                 with NamedTemporaryFile(dir='.', suffix=".pdf", delete=eval(os.getenv('DELETE_TEMP_PDF', 'True'))) as pdf:
                     pdf.write(uploaded_file.getbuffer())
+
+                    # Load Documents
                     loader = PyPDFLoader(pdf.name)
                     pages = loader.load()
 
@@ -65,6 +65,7 @@ if uploaded_files:
                     chunk_overlap = 0
                     docs = pages
 
+                    # Split Documents
                     if chunk_option:
                         docs = [Document('\n'.join([page.page_content for page in pages]))]
                         docs[0].metadata = {'source': pages[0].metadata['source']}
@@ -77,23 +78,22 @@ if uploaded_files:
                     )
                     chunks = text_splitter.split_documents(docs)
 
-                    threads = []
-                    threads.append(CustomThread(get_entity, (chunks, 'gsd')))
-                    threads.append(CustomThread(get_entity, (chunks, 'summ')))
-                    threads.append(CustomThread(get_entity, (chunks, 'all')))
-                    threads.append(CustomThread(get_entity_one, [c.page_content for c in chunks[:1]]))
-                    threads.append(CustomThread(get_table, pdf.name))
-
-                    [t.start() for t in threads]
-                    [t.join() for t in threads]
-                    
-                    result_gsd = threads[0].result
-                    result_summ = threads[1].result
-                    result = threads[2].result
-                    result_one = threads[3].result
-                    res_gene, res_snp, res_dis = threads[4].result
-
-                    # Combine
+                    # Start extraction process in parallel
+                    process = Process(model)
+                    with ThreadPoolExecutor() as executor:
+                        result_gsd = executor.submit(process.get_entity, (chunks, 'gsd'))
+                        result_summ = executor.submit(process.get_entity, (chunks, 'summ'))
+                        result = executor.submit(process.get_entity, (chunks, 'all'))
+                        result_one = executor.submit(process.get_entity_one, [c.page_content for c in chunks[:1]])
+                        result_table = executor.submit(process.get_table, pdf.name)
+
+                        result_gsd = result_gsd.result()
+                        result_summ = result_summ.result()
+                        result = result.result()
+                        result_one = result_one.result()
+                        res_gene, res_snp, res_dis = result_table.result()
+
+                    # Combine Result
                     result['Genes'] = res_gene + result_gsd['Genes']
                     result['SNPs'] = res_snp + result_gsd['SNPs']
                     result['Diseases'] = res_dis + result_gsd['Diseases']
@@ -119,9 +119,12 @@ if uploaded_files:
 
                     dataframe = pd.DataFrame(result)
                     dataframe = dataframe[['Genes', 'SNPs', 'Diseases', 'Title', 'Authors', 'Publisher Name', 'Publication Year', 'Population', 'Sample Size', 'Study Methodology', 'Study Level', 'Conclusion']]
+                    dataframe = dataframe[dataframe['Genes'].astype(bool)].reset_index(drop=True)
                     dataframe.drop_duplicates(['Genes', 'SNPs'], inplace=True)
                     dataframe.reset_index(drop=True, inplace=True)
-                    cleaned_df, cleaned_llm_df = validate(dataframe)
+                    
+                    # Validate Result
+                    cleaned_df, cleaned_llm_df = process.validate(dataframe)
 
                     end_time = datetime.now()
                     st.write("Success in ", round((end_time.timestamp() - start_time.timestamp()) / 60, 2), "minutes")

process.py

@@ -6,12 +6,12 @@ from langchain.chains.combine_documents.reduce import ReduceDocumentsChain
 from langchain.chains.combine_documents.stuff import StuffDocumentsChain
 from langchain.chains.llm import LLMChain
 from langchain.prompts import PromptTemplate
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import ChatOpenAI
 from pdf2image import convert_from_path
 from prompt import prompt_entity_gsd_chunk, prompt_entity_gsd_combine, prompt_entity_summ_chunk, prompt_entity_summ_combine, prompt_entities_chunk, prompt_entities_combine, prompt_entity_one_chunk, prompt_table, prompt_validation
 from table_detector import detection_transform, device, model, ocr, outputs_to_objects
 
-import google.generativeai as genai
 import io
 import json
 import os
@@ -20,11 +20,6 @@ import re
 import torch
 
 load_dotenv()
-genai.configure(api_key=os.environ['GOOGLE_API_KEY'])
-
-llm = ChatOpenAI(temperature=0, model_name="gpt-4-turbo")
-llm_p = ChatOpenAI(temperature=0, model_name="llama-3-sonar-large-32k-chat", api_key=os.environ['PERPLEXITY_API_KEY'], base_url="https://api.perplexity.ai")
-llm_g = genai.GenerativeModel(model_name='gemini-1.5-pro-latest')
 
 prompts = {
     'gsd': [prompt_entity_gsd_chunk, prompt_entity_gsd_combine],
@@ -32,211 +27,221 @@ prompts = {
     'all': [prompt_entities_chunk, prompt_entities_combine]
 }
 
-def get_entity(data):
+class Process():
 
-    chunks, types = data
+    def __init__(self, llm):
 
-    map_template = prompts[types][0]
-    map_prompt = PromptTemplate.from_template(map_template)
-    map_chain = LLMChain(llm=llm, prompt=map_prompt)
+        if llm.startswith('gpt'):
+            self.llm = ChatOpenAI(temperature=0, model_name=llm)
+        elif llm.startswith('gemini'):
+            self.llm = ChatGoogleGenerativeAI(temperature=0, model=llm)
+        else:
+            self.llm = ChatOpenAI(temperature=0, model_name=llm, api_key=os.environ['PERPLEXITY_API_KEY'], base_url="https://api.perplexity.ai")
 
-    reduce_template = prompts[types][1]
-    reduce_prompt = PromptTemplate.from_template(reduce_template)
-    reduce_chain = LLMChain(llm=llm, prompt=reduce_prompt)
+    def get_entity(self, data):
 
-    combine_chain = StuffDocumentsChain(
-        llm_chain=reduce_chain, document_variable_name="doc_summaries"
-    )
+        chunks, types = data
 
-    reduce_documents_chain = ReduceDocumentsChain(
-        combine_documents_chain=combine_chain,
-        collapse_documents_chain=combine_chain,
-        token_max=100000,
-    )
+        map_template = prompts[types][0]
+        map_prompt = PromptTemplate.from_template(map_template)
+        map_chain = LLMChain(llm=self.llm, prompt=map_prompt)
 
-    map_reduce_chain = MapReduceDocumentsChain(
-        llm_chain=map_chain,
-        reduce_documents_chain=reduce_documents_chain,
-        document_variable_name="docs",
-        return_intermediate_steps=False,
-    )
+        reduce_template = prompts[types][1]
+        reduce_prompt = PromptTemplate.from_template(reduce_template)
+        reduce_chain = LLMChain(llm=self.llm, prompt=reduce_prompt)
 
-    result = map_reduce_chain.invoke(chunks)['output_text']
-    print(types)
-    print(result)
-    if types != 'summ':
-        result = re.findall('(\{[^}]+\})', result)[0]
-        return eval(result)
-    
-    return result
+        combine_chain = StuffDocumentsChain(
+            llm_chain=reduce_chain, document_variable_name="doc_summaries"
+        )
 
-def get_entity_one(chunks):
+        reduce_documents_chain = ReduceDocumentsChain(
+            combine_documents_chain=combine_chain,
+            collapse_documents_chain=combine_chain,
+            token_max=100000,
+        )
 
-    result = llm.invoke(prompt_entity_one_chunk.format(chunks)).content
+        map_reduce_chain = MapReduceDocumentsChain(
+            llm_chain=map_chain,
+            reduce_documents_chain=reduce_documents_chain,
+            document_variable_name="docs",
+            return_intermediate_steps=False,
+        )
 
-    print('One')
-    print(result)
-    result = re.findall('(\{[^}]+\})', result)[0]
+        result = map_reduce_chain.invoke(chunks)['output_text']
+        print(types)
+        print(result)
+        if types != 'summ':
+            result = re.findall('(\{[^}]+\})', result)[0]
+            return eval(result)
+        
+        return result
 
-    return eval(result)
+    def get_entity_one(self, chunks):
 
-def get_table(path):
+        result = self.llm.invoke(prompt_entity_one_chunk.format(chunks)).content
+
+        print('One')
+        print(result)
+        result = re.findall('(\{[^}]+\})', result)[0]
+
+        return eval(result)
 
-    start_time = datetime.now()
-    images = convert_from_path(path)
-    print('PDF to Image', round((datetime.now().timestamp() - start_time.timestamp()) / 60, 2), "minutes")
-    tables = []
+    def get_table(self, path):
 
-    # Loop pages
-    for image in images:
+        start_time = datetime.now()
+        images = convert_from_path(path)
+        print('PDF to Image', round((datetime.now().timestamp() - start_time.timestamp()) / 60, 2), "minutes")
+        tables = []
 
-        pixel_values = detection_transform(image).unsqueeze(0).to(device)
-        with torch.no_grad():
-            outputs = model(pixel_values)
+        # Loop pages
+        for image in images:
 
-        id2label = model.config.id2label
-        id2label[len(model.config.id2label)] = "no object"
-        detected_tables = outputs_to_objects(outputs, image.size, id2label)
+            pixel_values = detection_transform(image).unsqueeze(0).to(device)
+            with torch.no_grad():
+                outputs = model(pixel_values)
 
-        # Loop table in page (if any)
-        for idx in range(len(detected_tables)):
-            cropped_table = image.crop(detected_tables[idx]["bbox"])
-            if detected_tables[idx]["label"] == 'table rotated':
-                cropped_table = cropped_table.rotate(270, expand=True)
+            id2label = model.config.id2label
+            id2label[len(model.config.id2label)] = "no object"
+            detected_tables = outputs_to_objects(outputs, image.size, id2label)
 
-            # TODO: what is the perfect threshold?
-            if detected_tables[idx]['score'] > 0.9:
-                print(detected_tables[idx])
-                tables.append(cropped_table)
+            # Loop table in page (if any)
+            for idx in range(len(detected_tables)):
+                cropped_table = image.crop(detected_tables[idx]["bbox"])
+                if detected_tables[idx]["label"] == 'table rotated':
+                    cropped_table = cropped_table.rotate(270, expand=True)
 
-    print('Detect table from image', round((datetime.now().timestamp() - start_time.timestamp()) / 60, 2), "minutes")
-    genes = []
-    snps = []
-    diseases = []
+                # TODO: what is the perfect threshold?
+                if detected_tables[idx]['score'] > 0.9:
+                    print(detected_tables[idx])
+                    tables.append(cropped_table)
 
-    # Loop tables
-    for table in tables:
+        print('Detect table from image', round((datetime.now().timestamp() - start_time.timestamp()) / 60, 2), "minutes")
+        genes = []
+        snps = []
+        diseases = []
 
-        buffer = io.BytesIO()
-        table.save(buffer, format='PNG')
-        image = Image(buffer)
+        # Loop tables
+        for table in tables:
+
+            buffer = io.BytesIO()
+            table.save(buffer, format='PNG')
+            image = Image(buffer)
+            
+            # Extract to dataframe
+            extracted_tables = image.extract_tables(ocr=ocr, implicit_rows=True, borderless_tables=True, min_confidence=0)
+
+            if len(extracted_tables) == 0:
+                continue
+
+            # Combine multiple dataframe
+            df_table = extracted_tables[0].df
+            for extracted_table in extracted_tables[1:]:
+                df_table = pd.concat([df_table, extracted_table.df]).reset_index(drop=True)
+
+            df_table.loc[0] = df_table.loc[0].fillna('')
+
+            # Identify multiple rows (in dataframe) as one row (in image)
+            rows = []
+            indexes = []
+            for i in df_table.index:
+                if not df_table.loc[i].isna().any():
+                    if len(indexes) > 0:
+                        rows.append(indexes)
+                    indexes = []
+                indexes.append(i)
+            rows.append(indexes)
+
+            df_table_cleaned = pd.DataFrame(columns=df_table.columns)
+            for row in rows:
+                row_str = df_table.loc[row[0]]
+                for idx in row[1:]:
+                    row_str += ' ' + df_table.loc[idx].fillna('')
+                row_str = row_str.str.strip()
+                df_table_cleaned.loc[len(df_table_cleaned)] = row_str
+
+            # Ask LLM with JSON data
+            json_table = df_table_cleaned.to_json(orient='records')
+            str_json_table = json.dumps(json.loads(json_table), indent=2)
+
+            result = self.llm.invoke(prompt_table.format(str_json_table)).content
+            print('table')
+            print(result)
+            result = result[result.find('['):result.rfind(']')+1]
+            try:
+                result = eval(result)
+            except SyntaxError:
+                result = []
+
+            for res in result:
+                res_gene = res['Genes']
+                res_snp = res['SNPs']
+                res_disease = res['Diseases']
+
+                for snp in res_snp:
+                    genes.append(res_gene)
+                    snps.append(snp)
+                    diseases.append(res_disease)
+
+        print('OCR table to extract', round((datetime.now().timestamp() - start_time.timestamp()) / 60, 2), "minutes")
+        print(genes, snps, diseases)
         
-        # Extract to dataframe
-        extracted_tables = image.extract_tables(ocr=ocr, implicit_rows=True, borderless_tables=True, min_confidence=0)
-
-        if len(extracted_tables) == 0:
-            continue
-
-        # Combine multiple dataframe
-        df_table = extracted_tables[0].df
-        for extracted_table in extracted_tables[1:]:
-            df_table = pd.concat([df_table, extracted_table.df]).reset_index(drop=True)
-
-        df_table.loc[0] = df_table.loc[0].fillna('')
-
-        # Identify multiple rows (in dataframe) as one row (in image)
-        rows = []
-        indexes = []
-        for i in df_table.index:
-            if not df_table.loc[i].isna().any():
-                if len(indexes) > 0:
-                    rows.append(indexes)
-                indexes = []
-            indexes.append(i)
-        rows.append(indexes)
-
-        df_table_cleaned = pd.DataFrame(columns=df_table.columns)
-        for row in rows:
-            row_str = df_table.loc[row[0]]
-            for idx in row[1:]:
-                row_str += ' ' + df_table.loc[idx].fillna('')
-            row_str = row_str.str.strip()
-            df_table_cleaned.loc[len(df_table_cleaned)] = row_str
-
-        # Ask LLM with JSON data
-        json_table = df_table_cleaned.to_json(orient='records')
+        return genes, snps, diseases
+
+    def validate(self, df):
+
+        df = df.fillna('')
+        df['Genes'] = df['Genes'].str.upper()
+        df['SNPs'] = df['SNPs'].str.lower()
+
+        # Check if there is two gene names
+        sym = ['-', '/', '|']
+        for i in df.index:
+            gene = df.loc[i, 'Genes']
+            for s in sym:
+                if s in gene:
+                    genes = gene.split(s)
+                    df.loc[i + 0.5] = df.loc[i]
+                    df = df.sort_index().reset_index(drop=True)
+                    df.loc[i, 'Genes'], df.loc[i + 1, 'Genes'] = genes[0], genes[1]
+
+        # Check if there is SNPs without 'rs'
+        for i in df.index:
+            safe = True
+            snp = df.loc[i, 'SNPs']
+            if re.fullmatch('rs(\d)+|', snp):
+                pass
+            elif re.fullmatch('ts(\d)+', snp):
+                snp = 'r' + snp[1:]
+            elif re.fullmatch('s(\d)+', snp):
+                snp = 'r' + snp
+            elif re.fullmatch('(\d)+', snp):
+                snp = 'rs' + snp
+            else:
+                safe = False
+                df = df.drop(i)
+                
+            if safe:
+                df.loc[i, 'SNPs'] = snp
+
+        df.reset_index(drop=True, inplace=True)
+
+        # Validate genes and diseases with LLM
+        json_table = df[['Genes', 'SNPs', 'Diseases']].to_json(orient='records')
         str_json_table = json.dumps(json.loads(json_table), indent=2)
 
-        result = llm.invoke(prompt_table.format(str_json_table)).content
-        print('table')
+        result = self.llm.invoke(input=prompt_validation.format(str_json_table)).content
+        print('val')
         print(result)
+
         result = result[result.find('['):result.rfind(']')+1]
         try:
             result = eval(result)
         except SyntaxError:
             result = []
 
-        for res in result:
-            res_gene = res['Genes']
-            res_snp = res['SNPs']
-            res_disease = res['Diseases']
-
-            for snp in res_snp:
-                genes.append(res_gene)
-                snps.append(snp)
-                diseases.append(res_disease)
-
-    print('OCR table to extract', round((datetime.now().timestamp() - start_time.timestamp()) / 60, 2), "minutes")
-    print(genes, snps, diseases)
-    
-    return genes, snps, diseases
-
-def validate(df):
-
-    df = df[df['Genes'].notna()].reset_index(drop=True)
-    df = df.fillna('')
-    df['Genes'] = df['Genes'].str.upper()
-    df['SNPs'] = df['SNPs'].str.lower()
-
-    # Check if there is two gene names
-    sym = ['-', '/', '|']
-    for i in df.index:
-        gene = df.loc[i, 'Genes']
-        for s in sym:
-            if s in gene:
-                genes = gene.split(s)
-                df.loc[i + 0.5] = df.loc[i]
-                df = df.sort_index().reset_index(drop=True)
-                df.loc[i, 'Genes'], df.loc[i + 1, 'Genes'] = genes[0], genes[1]
-
-    # Check if there is SNPs without 'rs'
-    for i in df.index:
-        safe = True
-        snp = df.loc[i, 'SNPs']
-        if re.fullmatch('rs(\d)+|', snp):
-            pass
-        elif re.fullmatch('ts(\d)+', snp):
-            snp = 'r' + snp[1:]
-        elif re.fullmatch('s(\d)+', snp):
-            snp = 'r' + snp
-        elif re.fullmatch('(\d)+', snp):
-            snp = 'rs' + snp
-        else:
-            safe = False
-            df = df.drop(i)
-            
-        if safe:
-            df.loc[i, 'SNPs'] = snp
-
-    df.reset_index(drop=True, inplace=True)
-
-    # Validate genes and diseases with LLM
-    json_table = df[['Genes', 'SNPs', 'Diseases']].to_json(orient='records')
-    str_json_table = json.dumps(json.loads(json_table), indent=2)
-
-    result = llm_p.invoke(input=prompt_validation.format(str_json_table)).content
-    print('val')
-    print(result)
-
-    result = result[result.find('['):result.rfind(']')+1]
-    try:
-        result = eval(result)
-    except SyntaxError:
-        result = []
-
-    df_val = pd.DataFrame(result)
-    df_val = df_val.merge(df.head(1).drop(['Genes', 'SNPs', 'Diseases'], axis=1), 'cross')
+        df_val = pd.DataFrame(result)
+        df_val = df_val.merge(df.head(1).drop(['Genes', 'SNPs', 'Diseases'], axis=1), 'cross')
 
-    # TODO: How to validate genes and SNPs with ground truth?
+        # TODO: How to validate genes and SNPs with ground truth?
 
-    return df, df_val
+        return df, df_val