naughtondale commited on
Commit
5e848dd
β€’
1 Parent(s): 8a1b20d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +151 -0
app.py ADDED
@@ -0,0 +1,151 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any
2
+ import gradio as gr
3
+ from langchain.embeddings.openai import OpenAIEmbeddings
4
+ from langchain.vectorstores import Chroma
5
+
6
+ from langchain.chains import ConversationalRetrievalChain
7
+ from langchain.chat_models import ChatOpenAI
8
+
9
+ from langchain.document_loaders import PyPDFLoader
10
+
11
+ import fitz
12
+ from PIL import Image
13
+
14
+ import chromadb
15
+ import re
16
+ import uuid
17
+
18
+ enable_box = gr.Textbox.update(value = None, placeholder = 'Upload your OpenAI API key',interactive = True)
19
+ disable_box = gr.Textbox.update(value = 'OpenAI API key is Set', interactive = False)
20
+
21
+ def set_apikey(api_key: str):
22
+ app.OPENAI_API_KEY = api_key
23
+ return disable_box
24
+
25
+ def enable_api_box():
26
+ return enable_box
27
+
28
+ def add_text(history, text: str):
29
+ if not text:
30
+ raise gr.Error('enter text')
31
+ history = history + [(text,'')]
32
+ return history
33
+
34
+ class my_app:
35
+ def __init__(self, OPENAI_API_KEY: str = None ) -> None:
36
+ self.OPENAI_API_KEY: str = OPENAI_API_KEY
37
+ self.chain = None
38
+ self.chat_history: list = []
39
+ self.N: int = 0
40
+ self.count: int = 0
41
+
42
+ def __call__(self, file: str) -> Any:
43
+ if self.count==0:
44
+ self.chain = self.build_chain(file)
45
+ self.count+=1
46
+ return self.chain
47
+
48
+ def chroma_client(self):
49
+ #create a chroma client
50
+ client = chromadb.Client()
51
+ #create a collecyion
52
+ collection = client.get_or_create_collection(name="my-collection")
53
+ return client
54
+
55
+ def process_file(self,file: str):
56
+ loader = PyPDFLoader(file.name)
57
+ documents = loader.load()
58
+ pattern = r"/([^/]+)$"
59
+ match = re.search(pattern, file.name)
60
+ file_name = match.group(1)
61
+ return documents, file_name
62
+
63
+ def build_chain(self, file: str):
64
+ documents, file_name = self.process_file(file)
65
+ #Load embeddings model
66
+ embeddings = OpenAIEmbeddings(openai_api_key=self.OPENAI_API_KEY)
67
+ pdfsearch = Chroma.from_documents(documents, embeddings, collection_name= file_name,)
68
+ chain = ConversationalRetrievalChain.from_llm(
69
+ ChatOpenAI(temperature=0.0, openai_api_key=self.OPENAI_API_KEY),
70
+ retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),
71
+ return_source_documents=True,)
72
+ return chain
73
+
74
+
75
+ def get_response(history, query, file):
76
+ if not file:
77
+ raise gr.Error(message='Upload a PDF')
78
+ chain = app(file)
79
+ result = chain({"question": query, 'chat_history':app.chat_history},return_only_outputs=True)
80
+ app.chat_history += [(query, result["answer"])]
81
+ app.N = list(result['source_documents'][0])[1][1]['page']
82
+ for char in result['answer']:
83
+ history[-1][-1] += char
84
+ yield history,''
85
+
86
+ def render_file(file):
87
+ doc = fitz.open(file.name)
88
+ page = doc[app.N]
89
+ #Render the page as a PNG image with a resolution of 300 DPI
90
+ pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
91
+ image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
92
+ return image
93
+
94
+ def render_first(file):
95
+ doc = fitz.open(file.name)
96
+ page = doc[0]
97
+ #Render the page as a PNG image with a resolution of 300 DPI
98
+ pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
99
+ image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
100
+ return image,[]
101
+
102
+ app = my_app()
103
+ with gr.Blocks() as demo:
104
+ with gr.Column():
105
+ with gr.Row():
106
+ with gr.Column(scale=0.8):
107
+ api_key = gr.Textbox(placeholder='Enter OpenAI API key', show_label=False, interactive=True).style(container=False)
108
+ with gr.Column(scale=0.2):
109
+ change_api_key = gr.Button('Change Key')
110
+ with gr.Row():
111
+ chatbot = gr.Chatbot(value=[], elem_id='chatbot').style(height=650)
112
+ show_img = gr.Image(label='Upload PDF', tool='select' ).style(height=680)
113
+ with gr.Row():
114
+ with gr.Column(scale=0.60):
115
+ txt = gr.Textbox(
116
+ show_label=False,
117
+ placeholder="Enter text and press enter",
118
+ ).style(container=False)
119
+ with gr.Column(scale=0.20):
120
+ submit_btn = gr.Button('submit')
121
+ with gr.Column(scale=0.20):
122
+ btn = gr.UploadButton("πŸ“ upload a PDF", file_types=[".pdf"]).style()
123
+
124
+ api_key.submit(
125
+ fn=set_apikey,
126
+ inputs=[api_key],
127
+ outputs=[api_key,])
128
+ change_api_key.click(
129
+ fn= enable_api_box,
130
+ outputs=[api_key])
131
+ btn.upload(
132
+ fn=render_first,
133
+ inputs=[btn],
134
+ outputs=[show_img,chatbot],)
135
+
136
+ submit_btn.click(
137
+ fn=add_text,
138
+ inputs=[chatbot,txt],
139
+ outputs=[chatbot, ],
140
+ queue=False).success(
141
+ fn=get_response,
142
+ inputs = [chatbot, txt, btn],
143
+ outputs = [chatbot,txt]).success(
144
+ fn=render_file,
145
+ inputs = [btn],
146
+ outputs=[show_img]
147
+ )
148
+
149
+
150
+ demo.queue()
151
+ demo.launch()