File size: 1,970 Bytes
3d3d712 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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
|