Update app.py

app.py

@@ -1,4 +1,5 @@
 from openai import OpenAI
+from ocr_mathpix import extract_pdf_mathpix
 import gradio as gr
 import fitz  # PyMuPDF
 from PIL import Image
@@ -28,11 +29,20 @@ def cal_tokens(message_data):
 def del_references(lines):
     # 定义正则表达式模式
     patterns = [
-        (r'\*\{.{0,5}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)\\section\*\{Tables', r'\section*{Tables\n'),
-        (r'\*\{.{0,5}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*)', ''),
-        (r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)(Table|Tables)', r'Tables'),
-        (r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)# SUPPLEMENTARY', r'# SUPPLEMENTARY'),
-        (r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)\[\^0\]', r'[^0]'),
+        (
+        r'\*\{.{0,5}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)\\section\*\{Tables',
+        r'\section*{Tables\n'),
+        (r'\*\{.{0,5}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*)',
+         ''),
+        (
+        r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)(Table|Tables)',
+        r'Tables'),
+        (
+        r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)# SUPPLEMENTARY',
+        r'# SUPPLEMENTARY'),
+        (
+        r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*?)\[\^0\]',
+        r'[^0]'),
         (r'#.{0,15}(References|Reference|REFERENCES|LITERATURE CITED|Referencesand notes|Notes and references)(.*)', '')
     ]
 
@@ -68,12 +78,25 @@ def extract_pdf_pypdf(pdf_dir):
     return file_content
 
 
-def openai_api(messages):
+def extract_pdf_md(pdf_dir):
+    print(f"start convert pdf 2 md: {pdf_dir}")
+    try:
+        content = extract_pdf_mathpix(pdf_folder_dir=os.path.split(pdf_dir)[0], pdf_dir=os.path.split(pdf_dir)[1],
+                                      md_folder_dir=os.path.split(pdf_dir)[0])
+    except Exception as e:
+        print(f"Error opening PDF: {e}")
+        return None
+
+    return content
+
+
+def openai_api(messages, model="claude-3-5-sonnet-20240620", temperature=0.1):
+    print("use model:", model, "temperature:", temperature)
     try:
         completion = client.chat.completions.create(
-            model="claude-3-5-sonnet-20240620",
+            model=model,
             messages=messages,
-            temperature=0.1,
+            temperature=temperature,
             max_tokens=8192,
             stream=True
         )
@@ -85,10 +108,10 @@ def openai_api(messages):
         return None
 
 
-def openai_chat_2_step(prompt, file_content):
+def openai_chat_2_step(prompt, file_content, model, temperature):
     all_response = ""
-    for i in range(len(file_content)//123000 + 1):
-        text = file_content[i*123000:(i+1)*123000]
+    for i in range(len(file_content) // 123000 + 1):
+        text = file_content[i * 123000:(i + 1) * 123000]
         # step1: 拆分两部分，前半部分
         messages = [
             {
@@ -101,7 +124,7 @@ def openai_chat_2_step(prompt, file_content):
         tokens = cal_tokens(messages)
         print("step一: 抽取部分{}：".format(i))
         print("prompt tokens:", tokens)
-        response_2_content = openai_api(messages)
+        response_2_content = openai_api(messages, model, temperature)
         if response_2_content:
             all_response += response_2_content + "\n"
 
@@ -128,11 +151,11 @@ Please pay attention to the pipe format as shown in the example below. This form
     tokens = cal_tokens(messages)
     print("step二: 合并部分：")
     print("prompt tokens:", tokens)
-    response = openai_api(messages)
+    response = openai_api(messages, model, temperature)
     return response
 
 
-def predict(prompt, file_content):
+def predict(prompt, file_content, model="claude-3-5-sonnet-20240620", temperature=0.1):
     file_content = del_references(file_content)
 
     messages = [
@@ -151,9 +174,9 @@ def predict(prompt, file_content):
     print("prompt tokens:", tokens)
     # time.sleep(20) # claude 需要加这个
     if tokens > 128000:
-        extract_result = openai_chat_2_step(prompt, file_content)
+        extract_result = openai_chat_2_step(prompt, file_content, model, temperature)
     else:
-        extract_result = openai_api(messages)
+        extract_result = openai_api(messages, model, temperature)
 
     return extract_result or "Too many users. Please wait a moment!"
 
@@ -254,6 +277,7 @@ def search_data_golden_Enzyme(keyword, selected_column):
     df = load_csv(CSV_FILE_PATH_Golden_Benchmark_Enzyme)
     return search_data(df, keyword, selected_column)
 
+
 def search_data_golden_Ribozyme(keyword, selected_column):
     df = load_csv(CSV_FILE_PATH_Golden_Benchmark_Ribozyme)
     return search_data(df, keyword, selected_column)
@@ -272,20 +296,21 @@ with gr.Blocks(title="Automated Enzyme Kinetics Extractor") as demo:
                 <p>How to use:
                 <br><strong>1</strong>: Upload your PDF.
                 <br><strong>2</strong>: Click "View PDF" to preview it.
-                <br><strong>3</strong>: Click "Extract Text" to extract Text.
+                <br><strong>3</strong>: Click "Convert to Markdown(Mathpix)/Convert to Text(PyMuPDF)" to extract PDF to Text.
                 <br><strong>4</strong>: Enter your extraction prompt in the input box.
-                <br><strong>5</strong>: Click "Generate" to extract, and the extracted information will display below.
+                <br><strong>5</strong>: Click "Generate" to extract data, and the extracted information will display below.
                 </p>'''
             )
             file_input = gr.File(label="Upload your PDF", type="filepath")
             example = gr.Examples(examples=[["./sample.pdf"]], inputs=file_input)
             with gr.Row():
                 viewer_button = gr.Button("View PDF", variant="secondary")
-                extract_button = gr.Button("Extract Text", variant="primary")
+                with gr.Row():
+                    extract_button_md = gr.Button("Convert to Markdown(Mathpix)", variant="primary")
+                    extract_button_text = gr.Button("Convert to Text(PyMuPDF)", variant="primary")
 
             with gr.Row():
                 with gr.Column(scale=1):
-
                     file_out = gr.Gallery(label="PDF Viewer", columns=1, height="auto", object_fit="contain")
 
                 with gr.Column(scale=1):
@@ -301,8 +326,17 @@ with gr.Blocks(title="Automated Enzyme Kinetics Extractor") as demo:
                     )
 
             with gr.Column():
-                model_input = gr.Textbox(lines=7, value=en_1, placeholder='Enter your extraction prompt here', label='Input Prompt')
+                model_input = gr.Textbox(lines=7, value=en_1, placeholder='Enter your extraction prompt here',
+                                         label='Input Prompt')
                 exp = gr.Button("Example Prompt")
+                with gr.Row():
+                    # 模型选择下拉菜单
+                    model_choices = ["claude-3-5-sonnet-20240620", "gpt-4o-2024-08-06"]
+                    model_dropdown = gr.Dropdown(choices=model_choices, label="Select Model", value=model_choices[0])
+
+                    # 温度选择滑块
+                    temp_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature", value=0.1)
+
                 with gr.Row():
                     gen = gr.Button("Generate", variant="primary")
                     clr = gr.Button("Clear")
@@ -335,7 +369,8 @@ with gr.Blocks(title="Automated Enzyme Kinetics Extractor") as demo:
             search_output = gr.HTML(label="Search Results", min_height=1000, max_height=1000)
 
             # 设置搜索功能
-            search_button.click(fn=search_data_golden_Enzyme, inputs=[search_box, column_dropdown], outputs=search_output)
+            search_button.click(fn=search_data_golden_Enzyme, inputs=[search_box, column_dropdown],
+                                outputs=search_output)
 
             # 将回车事件绑定到搜索按钮
             search_box.submit(fn=search_data_golden_Enzyme, inputs=[search_box, column_dropdown], outputs=search_output)
@@ -369,10 +404,12 @@ with gr.Blocks(title="Automated Enzyme Kinetics Extractor") as demo:
             search_output = gr.HTML(label="Search Results", min_height=1000, max_height=1000)
 
             # 设置搜索功能
-            search_button.click(fn=search_data_golden_Ribozyme, inputs=[search_box, column_dropdown], outputs=search_output)
+            search_button.click(fn=search_data_golden_Ribozyme, inputs=[search_box, column_dropdown],
+                                outputs=search_output)
 
             # 将回车事件绑定到搜索按钮
-            search_box.submit(fn=search_data_golden_Ribozyme, inputs=[search_box, column_dropdown], outputs=search_output)
+            search_box.submit(fn=search_data_golden_Ribozyme, inputs=[search_box, column_dropdown],
+                              outputs=search_output)
 
             # 初始加载整个 CSV 表格
             initial_output = load_csv(CSV_FILE_PATH_Golden_Benchmark_Ribozyme)
@@ -415,14 +452,13 @@ with gr.Blocks(title="Automated Enzyme Kinetics Extractor") as demo:
             else:
                 search_output.value = initial_output.to_html(classes='data', index=False, header=True)
 
-       
 
-    extract_button.click(extract_pdf_pypdf, inputs=file_input, outputs=text_output)
+
+    extract_button_md.click(extract_pdf_md, inputs=file_input, outputs=text_output)
+    extract_button_text.click(extract_pdf_pypdf, inputs=file_input, outputs=text_output)
     exp.click(update_input, outputs=model_input)
-    gen.click(fn=predict, inputs=[model_input, text_output], outputs=outputs)
+    gen.click(fn=predict, inputs=[model_input, text_output, model_dropdown, temp_slider], outputs=outputs)
     clr.click(fn=lambda: [gr.update(value=""), gr.update(value="")], inputs=None, outputs=[model_input, outputs])
     viewer_button.click(display_pdf_images, inputs=file_input, outputs=file_out)
 
-
 demo.launch()
-