githubear/oceanbase_chat_pdf_docker

1import os2import gradio as gr3# vectordb4from langchain.vectorstores import OceanBase5# embeddings6from langchain_community.embeddings import JinaEmbeddings7# PDF8# from PyPDF2 import PdfReader9from langchain_community.document_loaders import PyPDFLoader10from langchain.text_splitter import RecursiveCharacterTextSplitter11# LLM12from langchain_community.llms import Tongyi13from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder14from langchain_core.messages import HumanMessage15from langchain_core.runnables import RunnablePassthrough16from langchain.chains.combine_documents import create_stuff_documents_chain17 18def load_text_chunks_with_pdf_path(pdf_path):19    loader = PyPDFLoader(pdf_path)20    pdf_pages = loader.load()21    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=0)22    docs = text_splitter.split_documents(pdf_pages)23    return [doc.page_content for doc in docs]24 25## create vectore store 26def get_oceanbase() -> OceanBase:27    connection_str = OceanBase.connection_string_from_db_params(28        host="localhost",29        port="2881",30        database="test",31        user="root",32        password="",33    )34    embeddings = JinaEmbeddings(35        jina_api_key=os.environ.get("JINA_AI_API", ""), model_name="jina-embeddings-v2-base-zh"36    )37    # create oceanbase38    oceanbase = OceanBase(39        connection_string=connection_str,40        embedding_function=embeddings,41        # pre_delete_collection=True,  # TODO42    )43    return oceanbase44 45def text_rag_chain(retriever):46    llm = Tongyi()47    SYSTEM_TEMPLATE = """48    根据下面给出的上下文回答用户的问题。49    如果下面的上下文中不包含与问题相关的任何信息，请不要编造内容，仅仅回复”我不知道“50 51    <context>52    {context}53    </context>54    """55    question_answering_prompt = ChatPromptTemplate.from_messages(56        [57            (58                "system",59                SYSTEM_TEMPLATE,60            ),61            MessagesPlaceholder(variable_name="messages"),62        ]63    )64    document_chain = create_stuff_documents_chain(llm, question_answering_prompt)65 66    def parse_retriever_input(params):67        return params["messages"][-1].content68    retrieval_chain = RunnablePassthrough.assign(69        context=parse_retriever_input | retriever,70    ).assign(71        answer=document_chain,72    )73    return retrieval_chain74 75def run_pipeline(oceanbase, user_question, pdf_path):76    if pdf_path != "":77        texts = load_text_chunks_with_pdf_path(pdf_path)78        oceanbase.add_texts(texts=texts)79 80    retriever = oceanbase.as_retriever(k=5)81    chain = text_rag_chain(retriever)82    response = chain.invoke(83        {84            "messages": [85                HumanMessage(content=user_question)86            ],87        }88    )89    return response["answer"]90 91def gradio_interface(pdf_path, user_question):92    oceanbase = get_oceanbase()93    response = run_pipeline(oceanbase, user_question, str(pdf_path) if pdf_path else "")94    return response95 96iface = gr.Interface(97    fn=gradio_interface,98    inputs=[99        gr.File(label="Upload PDF files", file_types=['.pdf']),100        gr.Textbox(lines=2, placeholder="Ask a question from the PDF files")101    ],102    outputs="text",103    title="Chat with PDF",104    description="Ask questions from uploaded PDF documents"105)106 107if __name__ == "__main__":108    iface.launch(server_name="0.0.0.0", server_port=7860)