Openai API + langchain 分析小型pdf文档

声明:该版代码在2024.08.23有效。

代码如下:

python 复制代码
from langchain_community.document_loaders import PyPDFLoader
import getpass
import os
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

class QA:
    """
    A class to handle question-answering tasks on a given PDF document.

    Attributes:
        question (str): The question to be answered about the PDF.
        pdf_path (str): Path to the PDF file.
        model_name (str): Name of the model used for analysis.
        docs (list): Loaded PDF documents.
        vecstore (Chroma): The vector store object for storing document embeddings.

    Methods:
        set_environ(): Set environment variables for the OpenAI API.
        load_file(): Load a PDF file using PyPDFLoader.
        split_and_store(): Split the PDF text and store embeddings using Chroma.
        retrieve_pdf(): Retrieve and answer questions based on the PDF content.
    """
    def __init__(self, question, pdf_path, model_name):
        """
        Initializes the QA object with provided question, PDF path, and model name.

        Parameters:
            question (str): The question to be answered about the PDF.
            pdf_path (str): Path to the PDF file.
            model_name (str): Name of the model used for analysis.
        """
        self.question = question
        self.pdf_path = pdf_path
        self.model_name = model_name
        self.docs = None
        self.vecstore = None

    def set_environ(self):
        """
        Sets the environment variables necessary for OpenAI API authentication.
        """
        os.environ['OPENAI_API_KEY'] = input("your api:")
        os.environ['OPENAI_PROXY'] = 'http://127.0.0.1:20171'

    def load_file(self):
        """
        Loads the PDF file specified by the pdf_path attribute using PyPDFLoader.
        """
        loader = PyPDFLoader(self.pdf_path)
        self.docs = loader.load()

    def split_and_store(self):
        """
        Splits the loaded PDF text into manageable chunks and stores the embeddings in a vector store.
        """
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(self.docs)
        self.vecstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())

    def retrieve_pdf(self):
        """
        Retrieves context from the vector store and generates an answer to the input question
        using a retrieval-augmented generation chain.
        """
        retriever = self.vecstore.as_retriever()
        llm = ChatOpenAI(model="gpt-4o")

        system_prompt = (
            "You are an assistant for question-answering tasks. "
            "Use the following pieces of retrieved context to answer "
            "the question. If you don't know the answer, say that you "
            "don't know. Use three sentences maximum and keep the "
            "answer concise."
            "\n\n"
            "{context}"
        )

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system_prompt),
                ("human", "{input}"),
            ]
        )

        question_answer_chain = create_stuff_documents_chain(llm, prompt)
        rag_chain = create_retrieval_chain(retriever, question_answer_chain)

        results = rag_chain.invoke({"input": self.question})

        print(results['answer'])

    def run(self):
        self.set_environ()
        self.load_file()
        self.split_and_store()
        self.retrieve_pdf()

def __main__():
    """
    Main function to execute the QA class functionality.

    Prompts user for input parameters, creates a QA object, and processes the specified PDF.
    """
    question = input("Your question:")
    pdf_path = input("Enter the path of the pdf file:")
    model_name = input("Enter the model name:")
    qa = QA(question, pdf_path, model_name)
    qa.run()

if __name__ == "__main__":
    __main__()
相关推荐
我爱一条柴ya几秒前
【AI大模型】线性回归:经典算法的深度解析与实战指南
人工智能·python·算法·ai·ai编程
赶紧去巡山26 分钟前
pyhton基础【23】面向对象进阶四
python
旷世奇才李先生1 小时前
PyCharm 安装使用教程
ide·python·pycharm
这里有鱼汤1 小时前
“对象”?对象你个头!——Python世界观彻底崩塌的一天
后端·python
尘浮7281 小时前
60天python训练计划----day59
开发语言·python
wh39331 小时前
使用Python将PDF转换成word、PPT
python·pdf·word
船长@Quant2 小时前
数学视频动画引擎Python库 -- Manim Voiceover 语音服务 Speech Services
python·数学·manim·动画引擎·语音旁白
好开心啊没烦恼3 小时前
Python 数据分析:计算,分组统计1,df.groupby()。听故事学知识点怎么这么容易?
开发语言·python·数据挖掘·数据分析·pandas
lljss20204 小时前
Python11中创建虚拟环境、安装 TensorFlow
开发语言·python·tensorflow
空中湖4 小时前
tensorflow武林志第二卷第九章:玄功九转
人工智能·python·tensorflow