File size: 9,960 Bytes
0f98abb
0be4b49
 
 
 
 
 
 
27179da
7164e51
5aeea08
 
e5251a6
 
 
426b438
 
f5c78ec
 
 
0be4b49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
426b438
0be4b49
 
 
 
 
 
 
 
 
 
 
 
 
426b438
0be4b49
 
 
426b438
0be4b49
 
 
 
 
 
 
426b438
 
0be4b49
 
426b438
0be4b49
 
 
 
426b438
0be4b49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
665f9ef
0be4b49
 
 
 
 
 
 
 
 
 
 
 
426b438
0be4b49
 
 
 
426b438
039bca9
41866c2
039bca9
 
41866c2
039bca9
41866c2
039bca9
426b438
41866c2
0be4b49
 
 
0f98abb
 
 
 
853fad2
0f98abb
 
0be4b49
6e8b587
 
5aeea08
 
 
6e8b587
 
 
426b438
 
 
6e8b587
 
 
 
5aeea08
 
426b438
5aeea08
 
 
 
 
426b438
 
bdcb7b5
 
 
 
5aeea08
6e8b587
 
 
 
5aeea08
6e8b587
5aeea08
6e8b587
ee12eb1
e76c62c
 
6e8b587
 
 
 
 
 
 
 
 
 
e9ba6d1
6e8b587
426b438
6e8b587
ee12eb1
e76c62c
80a626b
426b438
 
e5251a6
a0777f9
 
426b438
 
a0777f9
 
e5251a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
426b438
 
e5251a6
426b438
 
 
 
 
 
 
 
 
 
 
 
dff4535
 
 
 
f5c78ec
dff4535
e9ba6d1
f5c78ec
426b438
 
 
 
 
 
 
 
 
223634f
 
 
 
 
426b438
 
 
f5c78ec
 
41866c2
426b438
 
dff4535
 
 
41866c2
 
f5c78ec
 
41866c2
 
f5c78ec
 
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
import streamlit as st
import urllib.request
import fitz
import re
import numpy as np
import tensorflow_hub as hub
import openai
from sklearn.neighbors import NearestNeighbors
import os
import time
import csv
from io import StringIO
import pandas as pd
from io import BytesIO
import base64
import threading
from queue import Queue
import logging

logging.basicConfig(level=logging.INFO)

def download_pdf(url, output_path):
    urllib.request.urlretrieve(url, output_path)

def preprocess(text):
    text = text.replace('\n', ' ')
    text = re.sub('\s+', ' ', text)
    return text

def pdf_to_text(path, start_page=1, end_page=None):
    doc = fitz.open(path)
    total_pages = doc.page_count

    if end_page is None:
        end_page = total_pages

    text_list = []

    for i in range(start_page-1, end_page):
        text = doc.load_page(i).get_text("text")
        text = preprocess(text)
        text_list.append(text)

    doc.close()
    return text_list

def text_to_chunks(texts, word_length=150, start_page=1):
    text_toks = [t.split(' ') for t in texts]
    page_nums = []
    chunks = []

    for idx, words in enumerate(text_toks):
        for i in range(0, len(words), word_length):
            chunk = words[i:i+word_length]
            if (i+word_length) > len(words) and (len(chunk) < word_length) and (
                len(text_toks) != (idx+1)):
                text_toks[idx+1] = chunk + text_toks[idx+1]
                continue
            chunk = ' '.join(chunk).strip()
            chunk = f'[Page no. {idx+start_page}]' + ' ' + '"' + chunk + '"'
            chunks.append(chunk)
    return chunks

class SemanticSearch:

    def __init__(self):
        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')
        self.fitted = False

    def fit(self, data, batch=1000, n_neighbors=5):
        self.data = data
        self.embeddings = self.get_text_embedding(data, batch=batch)
        n_neighbors = min(n_neighbors, len(self.embeddings))
        self.nn = NearestNeighbors(n_neighbors=n_neighbors)
        self.nn.fit(self.embeddings)
        self.fitted = True

    def call(self, text, return_data=True):
        inp_emb = self.use([text])
        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]

        if return_data:
            return [self.data[i] for i in neighbors]
        else:
            return neighbors

    def get_text_embedding(self, texts, batch=1000):
        embeddings = []
        for i in range(0, len(texts), batch):
            text_batch = texts[i:(i+batch)]
            emb_batch = self.use(text_batch)
            embeddings.append(emb_batch)
        embeddings = np.vstack(embeddings)
        return embeddings

def load_recommender(path, start_page=1):
    global recommender
    texts = pdf_to_text(path, start_page=start_page)
    chunks = text_to_chunks(texts, start_page=start_page)
    recommender.fit(chunks)
    return 'Corpus Loaded.'

def generate_text(openAI_key,prompt, engine="text-davinci-003"):
    openai.api_key = openAI_key
    completions = openai.Completion.create(
        engine=engine,
        prompt=prompt,
        max_tokens=512,
        n=1,
        stop=None,
        temperature=0.7,
    )
    message = completions.choices[0].text
    return message

def generate_answer(question,openAI_key):
    topn_chunks = recommender.call(question)
    prompt = ""
    prompt += 'search results:\n\n'
    for c in topn_chunks:
        prompt += c + '\n\n'

    prompt += "Instructions: Compose a comprehensive reply to the query using the search results given. "\
              "Cite each reference using [Page Number] notation (every result has this number at the beginning). "\
              "Citation should be done at the end of each sentence. If the search results mention multiple subjects "\
              "with the same name, create separate answers for each. Only include information found in the results and "\
              "don't add any additional information. Make sure the answer is correct, and don't output false content. "\
              "If the text does not relate to the query, simply state 'Found Nothing'. Ignore outlier "\
              "search results that have nothing to do with the question. Only answer what is asked. The "\
              "answer should be short and concise. \n\nQuery: {question}\nAnswer: "

    prompt += f"Query: {question}\nAnswer: "
    answer = generate_text(openAI_key, prompt,"text-davinci-003")
    return answer

recommender = SemanticSearch()

st.title('PDF GPT')

description = """ PDF GPT allows you to chat with your PDF file using Universal Sentence Encoder and Open AI. The returned response can cite the page number in square brackets([]) where the information is located, adding credibility to the responses and helping to locate pertinent information quickly."""

st.markdown(description)

openAI_key = st.sidebar.text_input('API Key', value='sk-')

data_section = st.sidebar.text_area("Paste Data:")
paste_data = st.sidebar.button("Paste Data")

add_row = st.sidebar.button("Add row")
row_count = st.session_state.get("row_count", 1)

num_concurrent_calls = st.sidebar.number_input("Concurrent Calls:", min_value=1, max_value=2000, value=10, step=1)
generate_all = st.sidebar.button("Generate All")

if add_row:
    row_count += 1
    st.session_state.row_count = row_count

if paste_data:
    data = StringIO(data_section.strip())
    reader = csv.reader(data, delimiter='\t', quotechar='"') # Changed delimiter to '\t'
    urls_questions = [row for row in reader]

    row_count = len(urls_questions)
    st.session_state.row_count = row_count

    for i, url_question in enumerate(urls_questions): # Directly iterate over urls_questions
        if len(url_question) >= 2:
            st.session_state[f"url{i}"] = url_question[0]
            st.session_state[f"question{i}"] = url_question[1]
        else:
            st.error(f"Row {i+1} does not have enough columns.")

for i in range(row_count):
    col1, col2, col3, col4 = st.columns(4)

    with col1:
        url = st.text_input(f'PDF URL {i+1}', key=f'url{i}', value=st.session_state.get(f'url{i}', ''))
    with col2:
        question = st.text_input(f'Question {i+1}', key=f'question{i}', value=st.session_state.get(f'question{i}', ''))
    with col3:
        # Initialize session state for answer if not already done
        if f'session_answer{i}' not in st.session_state:
            st.session_state[f'session_answer{i}'] = ''
    with col4:
        if st.button(f'Submit {i+1}'):
            if openAI_key.strip()=='':
                st.error('Please enter you Open AI Key')
            elif url.strip() == '':
                st.error('URL field is empty')
            elif question.strip() == '':
                st.error('Question field is empty')
            else:
                glob_url = url
                download_pdf(glob_url, 'corpus.pdf')
                load_recommender('corpus.pdf')

                answer = generate_answer(question,openAI_key)
                # Store the answer in session state
                st.session_state[f'session_answer{i}'] = answer
    with col3:
        answer_placeholder = st.empty()
        answer_placeholder.text_area(f'Answer {i+1}', key=f'answer{i}', value=st.session_state[f'session_answer{i}'])

def get_table_download_link(df, filename="data.csv", text="Download CSV file"):
    csv = df.to_csv(index=False)
    b64 = base64.b64encode(csv.encode()).decode() # some strings <-> bytes conversions necessary here
    href = f'{text}'
    return href

# Create a list of lists containing all URLs, questions, and answers
data = [[st.session_state.get(f'url{i}', ''), st.session_state.get(f'question{i}', ''), st.session_state.get(f'session_answer{i}', '')] for i in range(row_count)]

# Convert the data to a Pandas DataFrame
df = pd.DataFrame(data, columns=['URL', 'Question', 'Answer'])

# Generate a download link for the DataFrame
st.markdown(get_table_download_link(df), unsafe_allow_html=True)

def to_csv(data):
    output = BytesIO()
    writer = csv.writer(output)
    writer.writerows(data)
    return output.getvalue().decode('utf-8')

def get_table_download_link(df, filename="data.csv", text="Download CSV file"):
    csv = df.to_csv(index=False)
    b64 = base64.b64encode(csv.encode()).decode() # some strings <-> bytes conversions necessary here
    href = f'{text}'
    return href

class WorkerThread(threading.Thread):
    def __init__(self, jobs, results):
        super().__init__()
        self.jobs = jobs
        self.results = results

    def run(self):
        while True:
            job = self.jobs.get()
            if job is None:
                break
            try:
                i, question = job
                result = generate_answer(question, openAI_key)
                self.results.put((i, result))
                logging.info(f"Job {i} completed successfully.")
            except Exception as e:
                self.results.put((i, str(e)))
                logging.error(f"Error on job {i}: {str(e)}")

if generate_all:
    questions = [st.session_state.get(f"question{i}", "") for i in range(row_count)]

    jobs = Queue()
    results = Queue()

    workers = [WorkerThread(jobs, results) for _ in range(num_concurrent_calls)]

    # Start all worker threads
    for worker in workers:
        worker.start()

    # Join all worker threads
    for worker in workers:
        worker.join()

    logging.info("All worker threads have finished.")

    answers = {}
    while not results.empty():
        i, answer = results.get()
        if isinstance(answer, str) and 'Error' in answer:
            st.error(f"Error on row {i}: {answer}")
        else:
            answers[i] = answer

    logging.info(f"Collected {len(answers)} answers.")

    for i, answer in answers.items():
        st.session_state[f'session_answer{i}'] = answer

    logging.info("Session state updated with answers.")