File size: 15,643 Bytes
8bbef17
 
b657107
8bbef17
 
8705301
 
8bbef17
 
 
8705301
8bbef17
 
 
 
 
 
 
 
 
 
8705301
 
8bbef17
 
6e1a1ed
b657107
 
8bbef17
 
8705301
 
8bbef17
8705301
8bbef17
8705301
8bbef17
 
 
 
 
 
 
 
 
 
 
 
 
8705301
 
 
 
8bbef17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8705301
8bbef17
8705301
8bbef17
 
 
 
8705301
8bbef17
 
8705301
8bbef17
8705301
8bbef17
 
 
 
 
 
 
 
 
 
 
 
 
 
8705301
 
 
 
8bbef17
8705301
 
 
8bbef17
 
8705301
8bbef17
8705301
8bbef17
 
 
8705301
8bbef17
8705301
 
 
 
 
 
 
 
8bbef17
 
 
 
8705301
8bbef17
 
 
 
 
 
 
8705301
8bbef17
8705301
8bbef17
 
8705301
8bbef17
 
 
8705301
 
 
 
8bbef17
 
 
8705301
 
 
 
 
 
 
8bbef17
8705301
 
 
 
 
 
8bbef17
8705301
8bbef17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8705301
 
 
 
 
 
 
 
8bbef17
8705301
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8bbef17
8705301
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8bbef17
 
8705301
 
 
 
 
 
 
 
 
8bbef17
8705301
 
8bbef17
8705301
 
 
 
 
 
8bbef17
8705301
 
8bbef17
8705301
 
8bbef17
8705301
 
8bbef17
8705301
 
8bbef17
8705301
8bbef17
8705301
 
 
 
 
8bbef17
8705301
 
 
8bbef17
8705301
 
 
8bbef17
 
8705301
 
 
 
 
8bbef17
8705301
 
 
8bbef17
 
8705301
 
 
 
8bbef17
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
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379




import os
import io
import base64
import streamlit as st
import numpy as np
import fitz  # PyMuPDF
import tempfile
from ultralytics import YOLO
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity
from langchain_core.output_parsers import StrOutputParser
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import re
from PIL import Image
from streamlit_chat import message

# Load the trained model

model = YOLO("best.pt")
openai_api_key = os.environ.get("openai_api_key")

# Define the class indices for figures, tables, and text
figure_class_index = 4
table_class_index = 3

# Utility functions
def clean_text(text):
    return re.sub(r'\s+', ' ', text).strip()

def remove_references(text):
    reference_patterns = [
        r'\bReferences\b', r'\breferences\b', r'\bBibliography\b', r'\bCitations\b',
        r'\bWorks Cited\b', r'\bReference\b', r'\breference\b'
    ]
    lines = text.split('\n')
    for i, line in enumerate(lines):
        if any(re.search(pattern, line, re.IGNORECASE) for pattern in reference_patterns):
            return '\n'.join(lines[:i])
    return text

def save_uploaded_file(uploaded_file):
    temp_file = tempfile.NamedTemporaryFile(delete=False)
    temp_file.write(uploaded_file.getbuffer())
    temp_file.close()
    return temp_file.name

def summarize_pdf(pdf_file_path, num_clusters=10):
    embeddings_model = OpenAIEmbeddings(model="text-embedding-3-small", api_key=openai_api_key)
    llm = ChatOpenAI(model="gpt-3.5-turbo", api_key=openai_api_key, temperature=0.3)
    prompt = ChatPromptTemplate.from_template(
        """Could you please provide a concise and comprehensive summary of the given Contexts? 
        The summary should capture the main points and key details of the text while conveying the author's intended meaning accurately. 
        Please ensure that the summary is well-organized and easy to read, with clear headings and subheadings to guide the reader through each section. 
        The length of the summary should be appropriate to capture the main points and key details of the text, without including unnecessary information or becoming overly long. 
        example of summary:
        ## Summary:
        ## Key points:
        Contexts: {topic}"""
    )
    output_parser = StrOutputParser()
    chain = prompt | llm | output_parser

    loader = PyMuPDFLoader(pdf_file_path)
    docs = loader.load()
    full_text = "\n".join(doc.page_content for doc in docs)
    cleaned_full_text = clean_text(remove_references(full_text))

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0, separators=["\n\n", "\n", ".", " "])
    split_contents = text_splitter.split_text(cleaned_full_text)
    embeddings = embeddings_model.embed_documents(split_contents)

    kmeans = KMeans(n_clusters=num_clusters, init='k-means++', random_state=0).fit(embeddings)
    closest_point_indices = [np.argmin(np.linalg.norm(embeddings - center, axis=1)) for center in kmeans.cluster_centers_]
    extracted_contents = [split_contents[idx] for idx in closest_point_indices]

    results = chain.invoke({"topic": ' '.join(extracted_contents)})

    return generate_citations(results, extracted_contents)

def qa_pdf(pdf_file_path, query, num_clusters=5, similarity_threshold=0.6):
    embeddings_model = OpenAIEmbeddings(model="text-embedding-3-small", api_key=openai_api_key)
    llm = ChatOpenAI(model="gpt-3.5-turbo", api_key=openai_api_key, temperature=0.3)
    prompt = ChatPromptTemplate.from_template(
        """Please provide a detailed and accurate answer to the given question based on the provided contexts. 
        Ensure that the answer is comprehensive and directly addresses the query. 
        If necessary, include relevant examples or details from the text.
        Question: {question}
        Contexts: {contexts}"""
    )
    output_parser = StrOutputParser()
    chain = prompt | llm | output_parser

    loader = PyMuPDFLoader(pdf_file_path)
    docs = loader.load()
    full_text = "\n".join(doc.page_content for doc in docs)
    cleaned_full_text = clean_text(remove_references(full_text))

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=0, separators=["\n\n", "\n", ".", " "])
    split_contents = text_splitter.split_text(cleaned_full_text)
    embeddings = embeddings_model.embed_documents(split_contents)

    query_embedding = embeddings_model.embed_query(query)
    similarity_scores = cosine_similarity([query_embedding], embeddings)[0]
    top_indices = np.argsort(similarity_scores)[-num_clusters:]
    relevant_contents = [split_contents[i] for i in top_indices]

    results = chain.invoke({"question": query, "contexts": ' '.join(relevant_contents)})

    return generate_citations(results, relevant_contents, similarity_threshold)

def generate_citations(text, contents, similarity_threshold=0.6):
    embeddings_model = OpenAIEmbeddings(model="text-embedding-3-small", api_key=openai_api_key)
    text_sentences = re.split(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s', text)
    text_embeddings = embeddings_model.embed_documents(text_sentences)
    content_embeddings = embeddings_model.embed_documents(contents)
    similarity_matrix = cosine_similarity(text_embeddings, content_embeddings)

    cited_text = text
    relevant_sources = []
    source_mapping = {}
    sentence_to_source = {}

    for i, sentence in enumerate(text_sentences):
        if sentence in sentence_to_source:
            continue
        max_similarity = max(similarity_matrix[i])
        if max_similarity >= similarity_threshold:
            most_similar_idx = np.argmax(similarity_matrix[i])
            if most_similar_idx not in source_mapping:
                source_mapping[most_similar_idx] = len(relevant_sources) + 1
                relevant_sources.append((most_similar_idx, contents[most_similar_idx]))
            citation_idx = source_mapping[most_similar_idx]
            citation = f"([Source {citation_idx}](#source-{citation_idx}))"
            cited_sentence = re.sub(r'([.!?])$', f" {citation}\\1", sentence)
            sentence_to_source[sentence] = citation_idx
            cited_text = cited_text.replace(sentence, cited_sentence)

    sources_list = "\n\n## Sources:\n"
    for idx, (original_idx, content) in enumerate(relevant_sources):
        sources_list +=  f"""
<details style="margin: 1px 0; padding: 5px; border: 1px solid #ccc; border-radius: 8px; background-color: #f9f9f9; transition: all 0.3s ease;">
  <summary style="font-weight: bold; cursor: pointer; outline: none; padding: 5px 0; transition: color 0.3s ease;">Source {idx + 1}</summary>
  <pre style="white-space: pre-wrap; word-wrap: break-word; margin: 1px 0; padding: 10px; background-color: #fff; border-radius: 5px; border: 1px solid #ddd; box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);">{content}</pre>
</details>
"""

    # Add dummy blanks after the last source
    dummy_blanks = """
<div style="margin: 20px 0;"></div>
<div style="margin: 20px 0;"></div>
<div style="margin: 20px 0;"></div>
<div style="margin: 20px 0;"></div>
<div style="margin: 20px 0;"></div>

"""

    cited_text += sources_list + dummy_blanks
    return cited_text

def infer_image_and_get_boxes(image, confidence_threshold=0.8):
    results = model.predict(image)
    return [
        (int(box.xyxy[0][0]), int(box.xyxy[0][1]), int(box.xyxy[0][2]), int(box.xyxy[0][3]), int(box.cls[0]))
        for result in results for box in result.boxes
        if int(box.cls[0]) in {figure_class_index, table_class_index} and box.conf[0] > confidence_threshold
    ]

def crop_images_from_boxes(image, boxes, scale_factor):
    figures = []
    tables = []
    for (x1, y1, x2, y2, cls) in boxes:
        cropped_img = image[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)]
        if cls == figure_class_index:
            figures.append(cropped_img)
        elif cls == table_class_index:
            tables.append(cropped_img)
    return figures, tables

def process_pdf(pdf_file_path):
    doc = fitz.open(pdf_file_path)
    all_figures = []
    all_tables = []
    low_dpi = 50
    high_dpi = 300
    scale_factor = high_dpi / low_dpi
    low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
    
    for page_num, low_res_pix in enumerate(low_res_pixmaps):
        low_res_img = np.frombuffer(low_res_pix.samples, dtype=np.uint8).reshape(low_res_pix.height, low_res_pix.width, 3)
        boxes = infer_image_and_get_boxes(low_res_img)
        
        if boxes:
            high_res_pix = doc[page_num].get_pixmap(dpi=high_dpi)
            high_res_img = np.frombuffer(high_res_pix.samples, dtype=np.uint8).reshape(high_res_pix.height, high_res_pix.width, 3)
            figures, tables = crop_images_from_boxes(high_res_img, boxes, scale_factor)
            all_figures.extend(figures)
            all_tables.extend(tables)
    
    return all_figures, all_tables

def image_to_base64(img):
    buffered = io.BytesIO()
    img = Image.fromarray(img)
    img.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def on_btn_click():
    del st.session_state.chat_history[:]

# Streamlit interface

# Custom CSS for the file uploader
uploadercss='''
<style>
    [data-testid='stFileUploader'] {
        width: max-content;
    }
    [data-testid='stFileUploader'] section {
        padding: 0;
        float: left;
    }
    [data-testid='stFileUploader'] section > input + div {
        display: none;
    }
    [data-testid='stFileUploader'] section + div {
        float: right;
        padding-top: 0;
    }

</style>
'''

st.set_page_config(page_title="PDF Reading Assistant", page_icon="πŸ“„")

# Initialize chat history in session state if not already present
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

st.title("πŸ“„ PDF Reading Assistant")
st.markdown("### Extract tables, figures, summaries, and answers from your PDF files easily.")
chat_placeholder = st.empty()

# File uploader for PDF
uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
st.markdown(uploadercss, unsafe_allow_html=True)
if uploaded_file:
    file_path = save_uploaded_file(uploaded_file)

    # Chat container where all messages will be displayed
    chat_container = st.container()
    user_input = st.chat_input("Ask a question about the pdf......", key="user_input")
    with chat_container:
        # Scrollable chat messages
        for idx, chat in enumerate(st.session_state.chat_history):
            if chat.get("user"):
                message(chat["user"], is_user=True, allow_html=True, key=f"user_{idx}", avatar_style="initials", seed="user")
            if chat.get("bot"):
                message(chat["bot"], is_user=False, allow_html=True, key=f"bot_{idx}",seed="bot")

        # Input area and buttons for user interaction
        with st.form(key="chat_form", clear_on_submit=True,border=False):

            col1, col2, col3 = st.columns([1, 1, 1])
            with col1:
                summary_button = st.form_submit_button("Generate Summary")
            with col2:
                extract_button = st.form_submit_button("Extract Tables and Figures")
            with col3:
                st.form_submit_button("Clear message", on_click=on_btn_click)

            # Handle responses based on user input and button presses
            if summary_button:
                with st.spinner("Generating summary..."):
                    summary = summarize_pdf(file_path)
                st.session_state.chat_history.append({"user": "Generate Summary", "bot": summary})
                st.rerun()

            if extract_button:
                with st.spinner("Extracting tables and figures..."):
                    figures, tables = process_pdf(file_path)
                    if figures:
                        st.session_state.chat_history.append({"user": "Figures"})

                        for idx, figure in enumerate(figures):
                            figure_base64 = image_to_base64(figure)
                            result_html = f'<img src="data:image/png;base64,{figure_base64}" style="width:100%; display:block;" alt="Figure {idx+1}"/>'
                            st.session_state.chat_history.append({"bot": f"Figure {idx+1} {result_html}"})
                    if tables:
                        st.session_state.chat_history.append({"user": "Tables"})
                        for idx, table in enumerate(tables):
                            table_base64 = image_to_base64(table)
                            result_html = f'<img src="data:image/png;base64,{table_base64}" style="width:100%; display:block;" alt="Table {idx+1}"/>'
                            st.session_state.chat_history.append({"bot": f"Table {idx+1} {result_html}"})
                st.rerun()

            if user_input:
                st.session_state.chat_history.append({"user": user_input, "bot": None})
                with st.spinner("Processing..."):
                    answer = qa_pdf(file_path, user_input)
                st.session_state.chat_history[-1]["bot"] = answer
                st.rerun()

# Additional CSS and JavaScript to ensure the chat container is scrollable and scrolls to the bottom
st.markdown("""
    <style>
        #chat-container {
            max-height: 500px;
            overflow-y: auto;
            padding: 1rem;
            border: 1px solid #ddd;
            border-radius: 8px;
            background-color: #fefefe;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
            transition: background-color 0.3s ease;
        }
        #chat-container:hover {
            background-color: #f9f9f9;
        }
        .stChatMessage {
            padding: 0.75rem;
            margin: 0.75rem 0;
            border-radius: 8px;
            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
            transition: background-color 0.3s ease;
        }
        .stChatMessage--user {
            background-color: #E3F2FD;
        }
        .stChatMessage--user:hover {
            background-color: #BBDEFB;
        }
        .stChatMessage--bot {
            background-color: #EDE7F6;
        }
        .stChatMessage--bot:hover {
            background-color: #D1C4E9;
        }
        textarea {
            width: 100%;
            padding: 1rem;
            border: 1px solid #ddd;
            border-radius: 8px;
            box-shadow: inset 0 1px 3px rgba(0, 0, 0, 0.1);
            transition: border-color 0.3s ease, box-shadow 0.3s ease;
        }
        textarea:focus {
            border-color: #4CAF50;
            box-shadow: 0 0 5px rgba(76, 175, 80, 0.5);
        }
        .stButton > button {
            width: 100%;
            background-color: #4CAF50;
            color: white;
            border: none;
            border-radius: 8px;
            padding: 0.75rem;
            font-size: 16px;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
            transition: background-color 0.3s ease, box-shadow 0.3s ease;
        }
        .stButton > button:hover {
            background-color: #45A049;
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
        }
    </style>
    <script>
        const chatContainer = document.getElementById('chat-container');
        chatContainer.scrollTop = chatContainer.scrollHeight;
    </script>
""", unsafe_allow_html=True)