| | import os |
| |
|
| | from langchain.chat_models import AzureChatOpenAI, ChatOpenAI |
| | from langchain.document_loaders.pdf import PyPDFLoader |
| | from langchain.schema.messages import HumanMessage, SystemMessage |
| |
|
| | from taskweaver.plugin import Plugin, register_plugin |
| |
|
| | paper_summarize_prompt = r""" |
| | Please summarize this paper and highlight the key points, including the following: |
| | - The problem the paper is trying to solve. |
| | - The main idea of the paper. |
| | - The main contributions of the paper. |
| | - The main experiments and results of the paper. |
| | - The main conclusions of the paper. |
| | """ |
| |
|
| |
|
| | @register_plugin |
| | class SummarizePaperPlugin(Plugin): |
| | def __call__(self, paper_file_path: str): |
| | os.environ["OPENAI_API_TYPE"] = self.config.get("api_type", "azure") |
| | if os.environ["OPENAI_API_TYPE"] == "azure": |
| | model = AzureChatOpenAI( |
| | azure_endpoint=self.config.get("api_base"), |
| | openai_api_key=self.config.get("api_key"), |
| | openai_api_version=self.config.get("api_version"), |
| | azure_deployment=self.config.get("deployment_name"), |
| | temperature=0, |
| | verbose=True, |
| | ) |
| | elif os.environ["OPENAI_API_TYPE"] == "openai": |
| | os.environ["OPENAI_API_KEY"] = self.config.get("api_key") |
| | model = ChatOpenAI(model_name=self.config.get("deployment_name"), temperature=0, verbose=True) |
| | else: |
| | raise ValueError("Invalid API type. Please check your config file.") |
| |
|
| | loader = PyPDFLoader(paper_file_path) |
| | pages = loader.load() |
| |
|
| | messages = [ |
| | SystemMessage(content=paper_summarize_prompt), |
| | HumanMessage(content="The paper content:" + "\n".join([c.page_content for c in pages])), |
| | ] |
| |
|
| | summary_res = model.invoke(messages).content |
| |
|
| | description = f"We have summarized {len(pages)} pages of this paper." f"Paper summary is: {summary_res}" |
| |
|
| | return summary_res, description |
| |
|
