datak / app.py
LOUIS SANNA
clean(*): add black formatter
bddb702
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
from langchain.embeddings import OpenAIEmbeddings # for creating embeddings
from langchain.vectorstores import Chroma # for the vectorization part
from langchain.chains import ConversationalRetrievalChain
from langchain.llms import OpenAI # the LLM model we'll use (CHatGPT)
import gradio as gr
max_sources = 4
DB_DIR = "chroma"
embedding = OpenAIEmbeddings()
vectordb = Chroma(persist_directory=DB_DIR, embedding_function=embedding)
pdf_qa = ConversationalRetrievalChain.from_llm(
OpenAI(temperature=0.9, model_name="gpt-3.5-turbo"),
vectordb.as_retriever(),
return_source_documents=True,
)
def chat_pdf(query, chat_history=""):
result = pdf_qa({"question": query, "chat_history": chat_history})
answer = result["answer"]
source_docs = result["source_documents"]
print("source_docs", len(source_docs))
cleaned_docs = []
for doc in source_docs:
cleaned_content = doc.page_content
metadata_info = f"Metadata: {doc.metadata}\n"
cleaned_docs.append(metadata_info + cleaned_content)
# Pad the outputs to match the number of output components in the Gradio interface
padded_outputs = [answer] + cleaned_docs + [""] * (max_sources - len(cleaned_docs))
return padded_outputs
def create_outputs(num_sources):
outputs = [gr.outputs.Textbox(label="Answer")]
for i in range(1, num_sources + 1):
outputs.append(gr.outputs.Textbox(label=f"Source Document {i}"))
return outputs
iface = gr.Interface(
fn=chat_pdf,
inputs=[gr.inputs.Textbox(label="Query")],
outputs=create_outputs(max_sources),
examples=[
["Give 2 species of fulgoroidea"],
["What colors are found among fulgoroidea?"],
["Why are fulgoroidea so cute?"],
],
)
iface.launch(debug=True)