Rahatara commited on
Commit
bed15e6
·
verified ·
1 Parent(s): 1e639a8

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +94 -0
app.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import fitz # PyMuPDF
2
+ from sentence_transformers import SentenceTransformer
3
+ import numpy as np
4
+ import faiss
5
+ from typing import List, Dict
6
+
7
+ class MyApp:
8
+ def __init__(self) -> None:
9
+ self.documents = []
10
+ self.embeddings = None
11
+ self.index = None
12
+ self.model = SentenceTransformer('all-MiniLM-L6-v2')
13
+
14
+ def load_pdfs(self, file_paths: List[str]) -> None:
15
+ """Extracts text from multiple PDF files and stores them."""
16
+ self.documents = []
17
+ for file_path in file_paths:
18
+ doc = fitz.open(file_path)
19
+ for page_num in range(len(doc)):
20
+ page = doc[page_num]
21
+ text = page.get_text()
22
+ self.documents.append({"file": file_path, "page": page_num + 1, "content": text})
23
+ print("PDFs processed successfully!")
24
+
25
+ def build_vector_db(self) -> None:
26
+ """Builds a vector database using the content of the PDFs."""
27
+ if not self.documents:
28
+ print("No documents to process.")
29
+ return
30
+ contents = [doc["content"] for doc in self.documents]
31
+ self.embeddings = self.model.encode(contents, show_progress_bar=True)
32
+ self.index = faiss.IndexFlatL2(self.embeddings.shape[1])
33
+ self.index.add(np.array(self.embeddings))
34
+ print("Vector database built successfully!")
35
+
36
+ def search_documents(self, query: str, k: int = 3) -> List[Dict]:
37
+ """Searches for relevant document snippets using vector similarity."""
38
+ if not self.index:
39
+ print("Vector database is not built.")
40
+ return []
41
+ query_embedding = self.model.encode([query], show_progress_bar=False)
42
+ D, I = self.index.search(np.array(query_embedding), k)
43
+ results = [self.documents[i] for i in I[0]]
44
+ return results if results else [{"content": "No relevant documents found."}]
45
+
46
+ import gradio as gr
47
+ from typing import List, Tuple
48
+
49
+ app = MyApp()
50
+
51
+ def upload_files(files) -> str:
52
+ file_paths = [file.name for file in files]
53
+ app.load_pdfs(file_paths)
54
+ return f"Uploaded {len(files)} files successfully."
55
+
56
+ def build_vector_db() -> str:
57
+ app.build_vector_db()
58
+ return "Vector database built successfully!"
59
+
60
+ def respond(message: str, history: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], str]:
61
+ # Retrieve relevant documents
62
+ retrieved_docs = app.search_documents(message)
63
+ context = "\n".join([f"File: {doc['file']}, Page: {doc['page']}\n{doc['content']}" for doc in retrieved_docs])
64
+
65
+ # Generate response using the generative model
66
+ # Assuming you have set up the generative model as in your original code
67
+ # Replace the following line with your model's response generation
68
+ response_content = f"Simulated response based on the following context:\n{context}"
69
+
70
+ # Append the message and generated response to the chat history
71
+ history.append((message, response_content))
72
+ return history, ""
73
+
74
+ with gr.Blocks() as demo:
75
+ gr.Markdown("# PDF Chatbot")
76
+ gr.Markdown("Upload your PDFs, build a vector database, and start querying your documents.")
77
+
78
+ with gr.Row():
79
+ with gr.Column():
80
+ upload_btn = gr.File(label="Upload PDFs", file_types=[".pdf"], file_count="multiple")
81
+ upload_message = gr.Textbox(label="Upload Status", lines=2)
82
+ build_db_btn = gr.Button("Build Vector Database")
83
+ db_message = gr.Textbox(label="DB Build Status", lines=2)
84
+
85
+ upload_btn.change(upload_files, inputs=[upload_btn], outputs=[upload_message])
86
+ build_db_btn.click(build_vector_db, inputs=[], outputs=[db_message])
87
+
88
+ with gr.Column():
89
+ chatbot = gr.Chatbot(label="Chat Responses")
90
+ query_input = gr.Textbox(label="Enter your query here")
91
+ submit_btn = gr.Button("Submit")
92
+ submit_btn.click(respond, inputs=[query_input, chatbot], outputs=[chatbot, query_input])
93
+
94
+ demo.launch()