init

app.py

@@ -0,0 +1,103 @@
+import os
+import json
+import gradio as gr
+from typing import Literal
+
+TaskType = Literal["NER", "RE", "EE", "Base"]
+ModeType = Literal["direct", "fast-thinking", "slow-thinking", "customized"]
+
+
+class InterFace:
+    def __init__(self, dm):
+        """
+        receive the Pipeline instance as a parameter
+        """
+        self.dm = dm
+
+    def collect_and_process_inputs(self, task_description, file=None):
+        """
+        collect user inputs, return the processing result directly
+        """
+        # if file is not None:
+        #     file_path = self.save_file(file)
+        #     print(f"Your file has been saved, you can find it at: {file_path}")
+        # else:
+        #     file_path = None
+        file_path = "/disk/disk_20T/luoyujie/Agent/OneKE-Agent-reconstruct/data/input_files/Harry_Potter_Chapter1.pdf"
+        dm_response, dm_code = self.dm.get_extraction_result(file_path, task_description)
+        return json.dumps(dm_response, indent=2), dm_code
+
+    def save_file(self, file):
+        """
+        save file to the server and return the file path
+        """
+        file_path = os.path.join("saved", file.name)
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, "wb") as f:
+            f.write(file.read())
+        return file_path
+
+    def run_interface(self):
+        """
+        run Gradio interface
+        """
+        task_options = list(TaskType.__args__)
+        mode_options = list(ModeType.__args__)
+
+        # create the interface
+        with gr.Blocks() as demo:
+            gr.HTML("""
+                <div style="text-align:center;">
+                    <p align="center">
+                        <a href="https://github.com/R10836/DeepKE/blob/main/example/llm/assets/oneke_logo.png">
+                            <img src="https://raw.githubusercontent.com/R10836/DeepKE/refs/heads/main/example/llm/assets/oneke_logo.png" width="240"/>
+                        </a>
+                    </p>
+                    <h1>OneKE: A Dockerized Schema-Guided LLM Agent-based Knowledge Extraction System</h1>
+                    <p>
+                    📑[<a href="https://huggingface.co">Paper</a>]
+                    👨‍💻[<a href="https://github.com" target="_blank"><span class="icon"><i class="fab fa-github"></i></span>Code</a>]
+                    📄[<a href="https://zjunlp.gitbook.io">Docs</a>]
+                    🤗[<a href="https:/huggingface.cot" target="_blank">Demo</a>]
+                    </p>
+                </div>
+            """)
+
+            # section: input area
+            with gr.Row():
+                task = gr.Dropdown(choices=task_options, label="Select your Task", value="Base")
+                mode = gr.Dropdown(choices=mode_options, label="Select your Mode", value="direct")
+
+            task_description = gr.Textbox(label="Instruction", placeholder="Enter task description, e.g., extract key information from this article.")
+
+            text = gr.Textbox(label="Text", placeholder="Or enter text here directly")
+
+            use_file = gr.Checkbox(label="Use File", value=False)
+            file = gr.File(label="Upload File", visible=False, container=True, scale=2)
+            use_file.change(
+                lambda use_file: gr.update(visible=use_file),
+                inputs=[use_file],
+                outputs=[file]
+            )
+
+            # section: output area
+            code_output = gr.Code(label="Generated Schema", language="python", lines=10, interactive=False)
+            model_output = gr.Code(label="Final Answer", language="json", lines=10, interactive=False)
+
+            run_button = gr.Button("Submit")
+            run_button.click(
+                self.collect_and_process_inputs,
+                inputs=[task_description, file],
+                outputs=[model_output, code_output],
+            )
+
+        # launch the interface
+        demo.launch(share=False)
+
+
+if __name__ == "__main__":
+    # 你需要替换 dm 为你的 pipeline 实例
+    dm_instance = None  # 填入你的模型实例
+    app = InterFace(dm_instance)
+    app.run_interface()
+