|
|
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 |
|
|
|
