File size: 2,934 Bytes
56689a3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import gradio as gr
import os
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
# Create the pdfs directory if it doesn't exist
if not os.path.exists("pdfs"):
os.makedirs("pdfs")
def get_pdf_files():
"""Gets the list of PDF files from the 'pdfs' directory."""
return [f for f in os.listdir("pdfs") if f.endswith(".pdf")]
def index_pdfs():
"""Indexes the PDF files in the 'pdfs' directory."""
pdf_files = get_pdf_files()
if not pdf_files:
return "No PDF files found in the 'pdfs' directory."
success_files = []
failed_files = []
for pdf_file in pdf_files:
try:
file_path = os.path.join("pdfs", pdf_file)
if os.path.getsize(file_path) == 0:
failed_files.append(f"{pdf_file} (file is empty)")
continue
loader = PyPDFLoader(file_path)
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
embeddings = OpenAIEmbeddings()
db = Chroma.from_documents(texts, embeddings, persist_directory="./db")
db.persist()
success_files.append(pdf_file)
except Exception as e:
failed_files.append(f"{pdf_file} (Error: {e})")
status = ""
if success_files:
status += f"Successfully indexed: {', '.join(success_files)}\n"
if failed_files:
status += f"Failed to index: {', '.join(failed_files)}"
return status if status else "No files were processed."
def search(query):
"""Searches the indexed PDFs for the given query."""
embeddings = OpenAIEmbeddings()
db = Chroma(persist_directory="./db", embedding_function=embeddings)
docs = db.similarity_search(query)
results = ""
for doc in docs:
results += f"Source: {doc.metadata['source']}\n"
results += f"Content: {doc.page_content}\n\n"
return results
with gr.Blocks() as demo:
gr.Markdown("# Simple Semantic Search App")
with gr.Tab("Index PDFs"):
pdf_files_display = gr.Textbox(label="Available PDF Files", interactive=False, value="\n".join(get_pdf_files()))
index_button = gr.Button("Index PDFs")
index_status = gr.Textbox(label="Indexing Status", interactive=False)
index_button.click(index_pdfs, inputs=None, outputs=index_status)
with gr.Tab("Search"):
search_query = gr.Textbox(label="Search Query")
search_button = gr.Button("Search")
search_results = gr.Textbox(label="Search Results", interactive=False)
search_button.click(search, inputs=search_query, outputs=search_results)
if __name__ == "__main__":
demo.launch()
|