rag / app.py
dfasd
Update app.py
aaf8725 verified
raw
history blame
2.42 kB
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 environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# Initialize components
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'
# Load and process documents
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()
# Load prompt template
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()
)
# Gradio interface
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)