|
from dotenv import load_dotenv |
|
import os |
|
import gradio as gr |
|
from langchain_community.document_loaders import PyPDFLoader |
|
from langchain_text_splitters import CharacterTextSplitter |
|
from langchain_openai import OpenAIEmbeddings |
|
from langchain_community.vectorstores import Chroma |
|
from langchain_core.runnables import RunnablePassthrough |
|
from langchain_openai import ChatOpenAI |
|
from langchain import hub |
|
from langchain_core.output_parsers import StrOutputParser |
|
|
|
|
|
load_dotenv() |
|
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") |
|
|
|
|
|
text_splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=200, length_function=len) |
|
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY) |
|
llm = ChatOpenAI(model="gpt-4-1106-preview", api_key=OPENAI_API_KEY) |
|
vectordb_path = './vector_db' |
|
|
|
|
|
uploaded_files = ['airbus.pdf', 'annualreport2223.pdf'] |
|
dbname = 'vector_db' |
|
vectorstore = None |
|
|
|
for file in uploaded_files: |
|
loader = PyPDFLoader(file) |
|
data = loader.load() |
|
texts = text_splitter.split_documents(data) |
|
|
|
if vectorstore is None: |
|
vectorstore = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=os.path.join(vectordb_path, dbname)) |
|
else: |
|
vectorstore.add_documents(texts) |
|
|
|
vectorstore.persist() |
|
retriever = vectorstore.as_retriever() |
|
|
|
|
|
prompt = hub.pull("rlm/rag-prompt") |
|
print(prompt) |
|
|
|
def format_docs(docs): |
|
return "\n\n".join(doc.page_content for doc in docs) |
|
|
|
rag_chain = ( |
|
{"context": retriever | format_docs, "question": RunnablePassthrough()} |
|
| prompt |
|
| llm |
|
| StrOutputParser() |
|
) |
|
|
|
|
|
def rag_bot(query, chat_history): |
|
response = rag_chain.invoke({"input": query, "chat_history": chat_history}) |
|
return response |
|
|
|
chatbot = gr.Chatbot(avatar_images=["user.jpg", "bot.png"], height=600) |
|
clear_but = gr.Button(value="Clear Chat") |
|
|
|
def chat(query, chat_history): |
|
response = rag_bot(query, chat_history) |
|
chat_history.append((query, response)) |
|
return chat_history, chat_history |
|
|
|
demo = gr.Interface( |
|
fn=chat, |
|
inputs=["text", "state"], |
|
outputs=["chatbot", "state"], |
|
title="RAG Chatbot Prototype", |
|
description="A Chatbot using Retrieval-Augmented Generation (RAG) with PDF files.", |
|
allow_flagging="never", |
|
) |
|
|
|
if __name__ == '__main__': |
|
demo.launch(debug=True, share=True) |