ibrahim313 commited on
Commit
c8149d3
·
verified ·
1 Parent(s): ff4e25a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +127 -0
app.py ADDED
@@ -0,0 +1,127 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import base64
3
+ import gc
4
+ import tempfile
5
+ import uuid
6
+
7
+ import gradio as gr
8
+
9
+ # Importing necessary modules from llama_index
10
+ from llama_index.core import Settings
11
+ from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
12
+ from llama_index.llms.cohere import Cohere
13
+ from llama_index.embeddings.cohere import CohereEmbedding
14
+ from llama_index.postprocessor.cohere_rerank import CohereRerank
15
+ from llama_index.core import PromptTemplate
16
+
17
+ # Your Cohere API Key
18
+ API_KEY = "ziEpsRreaJzBi5HUDap7gMecJWXX69O26Hf71Kxo"
19
+
20
+ # Function to reset chat
21
+ def reset_chat():
22
+ gc.collect()
23
+
24
+ # Function to display PDF file
25
+ def display_pdf(file):
26
+ base64_pdf = base64.b64encode(file.read()).decode("utf-8")
27
+ pdf_display = f"""<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600px" type="application/pdf">
28
+ </iframe>"""
29
+ return pdf_display
30
+
31
+ # Function to process PDF and generate a query engine
32
+ def process_pdf(uploaded_file):
33
+ if uploaded_file:
34
+ try:
35
+ with tempfile.TemporaryDirectory() as temp_dir:
36
+ file_path = os.path.join(temp_dir, uploaded_file.name)
37
+ with open(file_path, "wb") as f:
38
+ f.write(uploaded_file.read())
39
+
40
+ # Creating an index over loaded data
41
+ loader = SimpleDirectoryReader(
42
+ input_dir=temp_dir,
43
+ required_exts=[".pdf"],
44
+ recursive=True
45
+ )
46
+ docs = loader.load_data()
47
+
48
+ # Setting up LLM & embedding model
49
+ llm = Cohere(api_key=API_KEY, model="command")
50
+ embed_model = CohereEmbedding(
51
+ cohere_api_key=API_KEY,
52
+ model_name="embed-english-v3.0",
53
+ input_type="search_query",
54
+ )
55
+
56
+ Settings.embed_model = embed_model
57
+ index = VectorStoreIndex.from_documents(docs, show_progress=True)
58
+
59
+ # Create a cohere reranker
60
+ cohere_rerank = CohereRerank(api_key=API_KEY)
61
+
62
+ # Create the query engine
63
+ Settings.llm = llm
64
+ query_engine = index.as_query_engine(streaming=True, node_postprocessors=[cohere_rerank])
65
+
66
+ # Customizing prompt template
67
+ qa_prompt_tmpl_str = (
68
+ "Context information is below.\n"
69
+ "---------------------\n"
70
+ "{context_str}\n"
71
+ "---------------------\n"
72
+ "Given the context information above, I want you to think step by step to answer the query in a crisp manner. If you don't know the answer, say 'I don't know!'.\n"
73
+ "Query: {query_str}\n"
74
+ "Answer: "
75
+ )
76
+ qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)
77
+
78
+ query_engine.update_prompts(
79
+ {"response_synthesizer:text_qa_template": qa_prompt_tmpl}
80
+ )
81
+
82
+ return query_engine, display_pdf(uploaded_file)
83
+ except Exception as e:
84
+ return None, f"An error occurred: {e}"
85
+ return None, "No file uploaded"
86
+
87
+ # Function to handle chat queries
88
+ def chat_with_pdf(prompt, query_engine):
89
+ if not query_engine:
90
+ return "Please upload and process a PDF file first."
91
+
92
+ full_response = ""
93
+ streaming_response = query_engine.query(prompt)
94
+
95
+ for chunk in streaming_response.response_gen:
96
+ full_response += chunk
97
+
98
+ return full_response
99
+
100
+ # Gradio Interface
101
+ with gr.Blocks() as demo:
102
+ gr.Markdown("# 🔍 Searchable Document Chatbot")
103
+ gr.Markdown("Upload your PDF document and start asking questions.")
104
+
105
+ pdf_file = gr.File(label="Upload your PDF file", file_types=[".pdf"])
106
+ pdf_preview = gr.HTML(label="PDF Preview")
107
+
108
+ query_engine = None
109
+ process_button = gr.Button("Process PDF")
110
+
111
+ chat_input = gr.Textbox(label="Ask a question")
112
+ chat_output = gr.Textbox(label="Chat Response")
113
+
114
+ def handle_pdf(file):
115
+ nonlocal query_engine
116
+ query_engine, pdf_html = process_pdf(file)
117
+ return pdf_html
118
+
119
+ def handle_chat(prompt):
120
+ return chat_with_pdf(prompt, query_engine)
121
+
122
+ process_button.click(fn=handle_pdf, inputs=pdf_file, outputs=pdf_preview)
123
+ chat_input.submit(fn=handle_chat, inputs=chat_input, outputs=chat_output)
124
+
125
+ gr.Markdown("Made with ❤️ by Muhammad Ibrahim Qasmi")
126
+
127
+ demo.launch()