Update app.py

app.py

@@ -1,88 +1,140 @@
 import gradio as gr
 import base64
-import os 
+import os
+from openai import OpenAI
+
 api_key = os.getenv('API_KEY')
+base_url = os.getenv("BASE_URL")
 
-def predict(input, file_input):
-    print("input:", input)
-    print("file_input:", file_input.name)
-    from gradio_client import Client
-    
-    client = Client(api_key)
-    extract_result = client.predict(
-        input,
-        file_input.name,
-        fn_index=1
-    )
-    if extract_result:
-        print(extract_result)
-        return extract_result
+client = OpenAI(
+    api_key=api_key,
+    base_url=base_url,
+)
+
+
+def extract_pdf_pypdf(pdf_dir):
+    import fitz
+    path = pdf_dir
+
+    try:
+        doc = fitz.open(path)
+    except:
+        print("can not read pdf")
+        return None
+
+    page_count = doc.page_count
+    file_content = ""
+    for page in range(page_count):
+        text = doc.load_page(page).get_text("text")
+        # 防止目录中包含References
+        file_content += text + "\n\n"
+
+    return file_content
+
+
+def openai_api(messages):
+    try:
+        completion = client.chat.completions.create(
+            model="claude-3-5-sonnet-20240620",
+            messages=messages,
+            temperature=0.1,
+            max_tokens=8192,
+            # timeout=300,
+            stream=True
+        )
+    except Exception as ex:
+        print("api 出现如下异常%s" % ex)
+        return None
+
+    if completion:
+        try:
+            response_2_list = [chunk.choices[0].delta.content if chunk.choices[0].delta.content else "" for chunk in
+                               completion]
+            print("response tokens:", len(response_2_list))
+
+            response_2_content = ''.join(response_2_list)
+            return response_2_content
+        except Exception as ex:
+            print("第二轮 出现如下异常%s" % ex)
+            return None
     else:
-        return "Too many user, please wait a monument!"
+        print("第二轮出现异常")
+        return None
+
+
+def predict(input_text, pdf_file):
+    if pdf_file is None:
+        return "Please upload a PDF file to proceed."
+
+    file_content = extract_pdf_pypdf(pdf_file.name)
+    messages = [
+        {
+            "role": "system",
+            "content": "You are an expert in information extraction from scientific literature.",
+        },
+        {"role": "user", "content": """Provided Text:
+    '''
+    {{""" + file_content + """}}
+    '''
+                                        """ + input_text}
+    ]
+    extract_result = openai_api(messages)
+
+    return extract_result or "Too many users. Please wait a moment!"
 
 
 def view_pdf(pdf_file):
+    if pdf_file is None:
+        return "Please upload a PDF file to view."
+
     with open(pdf_file.name, 'rb') as f:
         pdf_data = f.read()
-    # print("pdf_file", pdf_file)
-    # pdf_data = pdf_file
     b64_data = base64.b64encode(pdf_data).decode('utf-8')
-    # print("b64_data", b64_data)
     return f"<embed src='data:application/pdf;base64,{b64_data}' type='application/pdf' width='100%' height='700px' />"
 
 
-en_1 = ["""could you please help me extract the information of 'title'/'journal'/'year'/'author'/'institution'/'email' from the previous content in a markdown table format?
-If any of this information was not available in the paper, please replaced it with the string `""`, If the property contains multiple entities, please use a list to contain.
-"""]
+en_1 = """Could you please help me extract the information of 'title'/'journal'/'year'/'author'/'institution'/'email' from the previous content in a markdown table format?
+If any of this information was not available in the paper, please replace it with the string `""`. If the property contains multiple entities, please use a list to contain.
+"""
 
-en_2 = ["""could you please help me extract the information of 'title'/'journal'/'year'/'author'/'institution'/'email' from the previous content in a json format?
-If any of this information was not available in the paper, please replaced it with the string `""`, If the property contains multiple entities, please use a list to contain.
-"""]
+en_2 = """Could you please help me extract the information of 'title'/'journal'/'year'/'author'/'institution'/'email' from the previous content in a JSON format?
+If any of this information was not available in the paper, please replace it with the string `""`. If the property contains multiple entities, please use a list to contain.
+"""
 
-examples = [en_1, en_2]
+examples = [[en_1], [en_2]]
 
-with gr.Blocks(title="ChatPaperGPT") as demo:
+with gr.Blocks(title="PaperExtractGPT") as demo:
     gr.Markdown(
-        '''<p align="center" width="100%">
-        <img src="https://big-cheng.com/img/pdf.png" alt="pdf-logo" width="50"/>
-        <p>
-
+        '''<p align="center">
         <h1 align="center"> Paper Extract GPT </h1>
         <p> How to use:
-        <br> <strong>#1</strong>: Upload your pdf.
-        <br> <strong>#2</strong>: Click the View PDF button to view it.
-        <br> <strong>#3</strong>: Enter your extraction prompt in the input box (of course, you can click example to test).
-        <br> <strong>#4</strong>: Click Generate to extract, and the extracted information will be displayed in markdown form.
+        <br> <strong>1</strong>: Upload your PDF.
+        <br> <strong>2</strong>: Click "View PDF" to preview it.
+        <br> <strong>3</strong>: Enter your extraction prompt in the input box.
+        <br> <strong>4</strong>: Click "Generate" to extract, and the extracted information will display below.
         </p>
         '''
     )
     with gr.Row():
         with gr.Column():
             gr.Markdown('## Upload PDF')
-            file_input = gr.File(type="filepath")
+            file_input = gr.File(label="Upload your PDF", type="filepath")
             viewer_button = gr.Button("View PDF")
-            file_out = gr.HTML()
+            file_out = gr.HTML(label="PDF Preview")
+
         with gr.Column():
-            with gr.Row():
-                model_input = gr.Textbox(lines=7, placeholder='Input prompt about extract information from paper',
-                                         label='Input')
+            model_input = gr.Textbox(lines=7, placeholder='Enter your extraction prompt here', label='Input Prompt')
+            example = gr.Examples(examples=examples, inputs=model_input)
             with gr.Row():
                 gen = gr.Button("Generate")
                 clr = gr.Button("Clear")
-            example = gr.Examples(examples=examples, inputs=model_input)
-
-            with gr.Row():
-                outputs = gr.Markdown(label='Output', show_label=True, value="""| Title                                       | Journal            | Year | Author                                        | Institution                                           | Email                 |
+            outputs = gr.Markdown(label='Output', show_label=True,  value="""| Title                                       | Journal            | Year | Author                                        | Institution                                           | Email                 |
 |---------------------------------------------|--------------------|------|-----------------------------------------------|-------------------------------------------------------|-----------------------|
 | Paleomagnetic Study of Deccan Traps from Jabalpur to Amarkantak, Central India | J. Geomag. Geoelectr. | 1973 | R. K. VERMA, G. PULLAIAH, G.R. ANJANEYULU, P. K. MALLIK | National Geophysical Research Institute, Hyderabad, and Indian School o f Mines, Dhanbad | "" |
 """)
 
-    inputs = [model_input, file_input]
-    gen.click(fn=predict, inputs=inputs, outputs=outputs)
-    clr.click(fn=lambda value: [gr.update(value=""), gr.update(value="")], inputs=clr,
-              outputs=[model_input, outputs])
-
+    gen.click(fn=predict, inputs=[model_input, file_input], outputs=outputs)
+    clr.click(fn=lambda: [gr.update(value=""), gr.update(value="")], inputs=None, outputs=[model_input, outputs])
     viewer_button.click(view_pdf, inputs=file_input, outputs=file_out)
-    # parser_button.click(extract_text, inputs=file_input, outputs=[xml_out, md_out, rich_md_out])
 
 demo.launch()