aaporosh commited on
Commit
b561129
·
verified ·
1 Parent(s): 11694c7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +406 -146
app.py CHANGED
@@ -1,157 +1,417 @@
1
- # ------------- app.py -------------
2
  import streamlit as st
3
- from pathlib import Path
 
4
  from io import BytesIO
5
- import pdfplumber, pytesseract, time, re, logging, os
 
6
  from PIL import Image
7
- from langchain.text_splitter import RecursiveCharacterTextSplitter
8
  from langchain_community.vectorstores import FAISS
9
  from sentence_transformers import SentenceTransformer
10
- from transformers import pipeline
11
- import numpy as np
 
 
 
 
 
12
 
13
- logging.basicConfig(level=logging.INFO)
 
14
  logger = logging.getLogger(__name__)
15
 
16
- ###############################################################################
17
- # Page layout
18
- ###############################################################################
19
- st.set_page_config(page_title="PDF Chat & Summarize", layout="wide")
20
- st.markdown("""
21
- <style>
22
- .block-container { padding-top: 1rem; padding-bottom: 0; }
23
- .stTabs [data-baseweb="tab-list"] { gap: 4px; }
24
- .stTabs [data-baseweb="tab"] { padding: 8px 24px; }
25
- .chat-msg { padding: 0.5rem 1rem; border-radius: 8px; margin: 0.3rem 0; }
26
- .user { background-color: #e3f2fd; margin-left: 20%; }
27
- .assistant { background-color: #f1f3f4; margin-right: 20%; }
28
- </style>
29
- """, unsafe_allow_html=True)
30
-
31
- ###############################################################################
32
- # Cached heavy objects
33
- ###############################################################################
34
- @st.cache_resource(show_spinner=False)
35
- def load_embed():
36
- return SentenceTransformer("all-MiniLM-L6-v2")
37
-
38
- @st.cache_resource(show_spinner=False)
39
- def load_qa():
40
- return pipeline("text2text-generation", model="google/flan-t5-large", max_length=512)
41
-
42
- @st.cache_resource(show_spinner=False)
43
- def load_sum():
44
- return pipeline("summarization", model="facebook/bart-large-cnn", max_length=250)
45
-
46
- embed = load_embed()
47
- qa_pipe = load_qa()
48
- sum_pipe = load_sum()
49
-
50
- ###############################################################################
51
- # Helpers
52
- ###############################################################################
53
- def extract_pdf(uploaded_file):
54
- """Return (plain text, image_list)"""
55
- text = ""
56
- images = []
57
- with pdfplumber.open(BytesIO(uploaded_file.getbuffer())) as pdf:
58
- for page in pdf.pages:
59
- txt = page.extract_text_layout() or page.extract_text()
60
- if not txt:
61
- img = page.to_image(resolution=200).original
62
- txt = pytesseract.image_to_string(img)
63
- text += txt + "\n"
64
- for img in page.images:
65
- try:
66
- x0, y0, x1, y1 = img["x0"], img["y0"], img["x1"], img["y1"]
67
- pil = page.within_bbox((x0, y0, x1, y1)).to_image(resolution=200).original
68
- images.append(pil)
69
- except Exception:
70
- pass
71
- return text.strip(), images
72
-
73
- def build_index(text):
74
- splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=80)
75
- chunks = splitter.split_text(text)
76
- vectors = embed.encode(chunks, show_progress_bar=False, batch_size=64)
77
- index = FAISS.from_embeddings(list(zip(chunks, vectors)), embed)
78
- return index
79
-
80
- def summarize(text):
81
- if len(text) < 50:
82
- return "Document too short to summarize."
83
- # pick top 3k chars to stay within model limit
84
- truncated = text[:3000]
85
- return sum_pipe(truncated, max_length=250, min_length=60, do_sample=False)[0]["summary_text"]
86
-
87
- def answer(question, index):
88
- if index is None:
89
- return "Please upload & process a PDF first."
90
- docs = index.similarity_search(question, k=4)
91
- context = "\n".join([d.page_content for d in docs])
92
- prompt = f"Answer the question using ONLY the context below.\n\nContext:\n{context}\n\nQuestion: {question}"
93
- return qa_pipe(prompt, max_length=256, do_sample=False)[0]["generated_text"]
94
-
95
- ###############################################################################
96
- # Session init
97
- ###############################################################################
98
- if "messages" not in st.session_state:
99
- st.session_state.messages = []
100
- if "index" not in st.session_state:
101
- st.session_state.index = None
102
- if "raw_text" not in st.session_state:
103
- st.session_state.raw_text = ""
104
- if "images" not in st.session_state:
105
- st.session_state.images = []
106
-
107
- ###############################################################################
108
- # Sidebar
109
- ###############################################################################
110
- with st.sidebar:
111
- st.subheader("📁 Upload PDF")
112
- uploaded = st.file_uploader("Choose a file", type="pdf", label_visibility="collapsed")
113
- if uploaded and st.button("Process PDF"):
114
- with st.spinner("Extracting text & images…"):
115
- st.session_state.raw_text, st.session_state.images = extract_pdf(uploaded)
116
- st.session_state.index = build_index(st.session_state.raw_text)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
117
  st.session_state.messages = []
118
- st.toast("PDF ready!")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
119
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
120
  if st.session_state.images:
121
- st.subheader("🖼️ Extracted Images")
122
- for im in st.session_state.images:
123
- st.image(im, use_column_width=True)
124
-
125
- ###############################################################################
126
- # Main Tabs
127
- ###############################################################################
128
- tab_chat, tab_sum = st.tabs(["💬 Chat", "📄 Summarize"])
129
-
130
- with tab_chat:
131
- if st.session_state.index is None:
132
- st.info("Upload & process a PDF first using the sidebar.")
133
- else:
134
- # history
135
- for role, msg in st.session_state.messages:
136
- css = "user" if role == "user" else "assistant"
137
- st.markdown(f'<div class="chat-msg {css}">{msg}</div>', unsafe_allow_html=True)
138
-
139
- # input
140
- if question := st.chat_input("Ask anything about the PDF…"):
141
- st.session_state.messages.append(("user", question))
142
- st.markdown(f'<div class="chat-msg user">{question}</div>', unsafe_allow_html=True)
143
-
144
- with st.spinner("Thinking…"):
145
- resp = answer(question, st.session_state.index)
146
- st.session_state.messages.append(("assistant", resp))
147
- st.markdown(f'<div class="chat-msg assistant">{resp}</div>', unsafe_allow_html=True)
148
-
149
- with tab_sum:
150
- if not st.session_state.raw_text:
151
- st.info("Upload & process a PDF first.")
152
- else:
153
- if st.button("Generate Summary"):
154
- with st.spinner("Summarizing…"):
155
- summary = summarize(st.session_state.raw_text)
156
- st.subheader("Summary")
157
- st.write(summary)
 
 
1
  import streamlit as st
2
+ import logging
3
+ import os
4
  from io import BytesIO
5
+ import pdfplumber
6
+ from pdf2image import convert_from_bytes
7
  from PIL import Image
8
+ from langchain.text_splitter import CharacterTextSplitter
9
  from langchain_community.vectorstores import FAISS
10
  from sentence_transformers import SentenceTransformer
11
+ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments
12
+ from datasets import load_dataset
13
+ from rank_bm25 import BM25Okapi
14
+ from rouge_score import rouge_scorer
15
+ import re
16
+ import time
17
+ import pytesseract
18
 
19
+ # Setup logging for Spaces
20
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
21
  logger = logging.getLogger(__name__)
22
 
23
+ # Lazy load models
24
+ @st.cache_resource(ttl=1800)
25
+ def load_embeddings_model():
26
+ logger.info("Loading embeddings model")
27
+ try:
28
+ return SentenceTransformer("all-MiniLM-L6-v2")
29
+ except Exception as e:
30
+ logger.error(f"Embeddings load error: {str(e)}")
31
+ st.error(f"Embedding model error: {str(e)}")
32
+ return None
33
+
34
+ @st.cache_resource(ttl=1800)
35
+ def load_qa_pipeline():
36
+ logger.info("Loading QA pipeline")
37
+ try:
38
+ dataset = load_and_prepare_dataset()
39
+ if dataset:
40
+ fine_tuned_pipeline = fine_tune_qa_model(dataset)
41
+ if fine_tuned_pipeline:
42
+ return fine_tuned_pipeline
43
+ return pipeline("text2text-generation", model="google/flan-t5-small", max_length=300)
44
+ except Exception as e:
45
+ logger.error(f"QA model load error: {str(e)}")
46
+ st.error(f"QA model error: {str(e)}")
47
+ return None
48
+
49
+ @st.cache_resource(ttl=1800)
50
+ def load_summary_pipeline():
51
+ logger.info("Loading summary pipeline")
52
+ try:
53
+ return pipeline("summarization", model="facebook/bart-large-cnn", max_length=250)
54
+ except Exception as e:
55
+ logger.error(f"Summary model load error: {str(e)}")
56
+ st.error(f"Summary model error: {str(e)}")
57
+ return None
58
+
59
+ # Load and prepare dataset (e.g., SQuAD)
60
+ @st.cache_data(ttl=3600)
61
+ def load_and_prepare_dataset(dataset_name="squad", max_samples=1000):
62
+ logger.info(f"Loading dataset: {dataset_name}")
63
+ try:
64
+ dataset = load_dataset(dataset_name, split="train[:80%]")
65
+ dataset = dataset.shuffle(seed=42).select(range(min(max_samples, len(dataset))))
66
+
67
+ def preprocess(examples):
68
+ inputs = [f"question: {q} context: {c}" for q, c in zip(examples['question'], examples['context'])]
69
+ targets = examples['answers']['text']
70
+ return {'input_text': inputs, 'target_text': [t[0] if t else "" for t in targets]}
71
+
72
+ dataset = dataset.map(preprocess, batched=True, remove_columns=dataset.column_names)
73
+ return dataset
74
+ except Exception as e:
75
+ logger.error(f"Dataset load error: {str(e)}")
76
+ return None
77
+
78
+ # Fine-tune QA model
79
+ @st.cache_resource(ttl=3600)
80
+ def fine_tune_qa_model(dataset):
81
+ logger.info("Starting fine-tuning")
82
+ try:
83
+ model_name = "google/flan-t5-small"
84
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
85
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
86
+
87
+ def tokenize_function(examples):
88
+ model_inputs = tokenizer(examples['input_text'], max_length=512, truncation=True, padding="max_length")
89
+ labels = tokenizer(examples['target_text'], max_length=128, truncation=True, padding="max_length")
90
+ model_inputs["labels"] = labels["input_ids"]
91
+ return model_inputs
92
+
93
+ tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=['input_text', 'target_text'])
94
+
95
+ training_args = TrainingArguments(
96
+ output_dir="./fine_tuned_model",
97
+ num_train_epochs=2,
98
+ per_device_train_batch_size=4,
99
+ save_steps=500,
100
+ logging_steps=100,
101
+ evaluation_strategy="no",
102
+ learning_rate=3e-5,
103
+ fp16=False,
104
+ )
105
+
106
+ trainer = Trainer(
107
+ model=model,
108
+ args=training_args,
109
+ train_dataset=tokenized_dataset,
110
+ )
111
+ trainer.train()
112
+
113
+ model.save_pretrained("./fine_tuned_model")
114
+ tokenizer.save_pretrained("./fine_tuned_model")
115
+ logger.info("Fine-tuning complete")
116
+ return pipeline("text2text-generation", model="./fine_tuned_model", tokenizer="./fine_tuned_model", max_length=300)
117
+ except Exception as e:
118
+ logger.error(f"Fine-tuning error: {str(e)}")
119
+ return None
120
+
121
+ # Augment vector store with dataset
122
+ def augment_vector_store(vector_store, dataset_name="squad", max_samples=300):
123
+ logger.info(f"Augmenting vector store with dataset: {dataset_name}")
124
+ try:
125
+ dataset = load_dataset(dataset_name, split="train").select(range(min(max_samples, len(dataset))))
126
+ chunks = [f"Context: {c}\nAnswer: {a['text'][0]}" for c, a in zip(dataset['context'], dataset['answers'])]
127
+ embeddings_model = load_embeddings_model()
128
+ if embeddings_model and vector_store:
129
+ embeddings = embeddings_model.encode(chunks, batch_size=128, show_progress_bar=False)
130
+ vector_store.add_embeddings(zip(chunks, embeddings))
131
+ return vector_store
132
+ except Exception as e:
133
+ logger.error(f"Vector store augmentation error: {str(e)}")
134
+ return vector_store
135
+
136
+ # Process PDF with enhanced extraction and OCR fallback
137
+ def process_pdf(uploaded_file):
138
+ logger.info("Processing PDF with enhanced extraction")
139
+ try:
140
+ text = ""
141
+ code_blocks = []
142
+ images = []
143
+ with pdfplumber.open(BytesIO(uploaded_file.getvalue())) as pdf:
144
+ for page in pdf.pages[:8]:
145
+ extracted = page.extract_text(layout=False)
146
+ if not extracted:
147
+ try:
148
+ img = page.to_image(resolution=150).original
149
+ extracted = pytesseract.image_to_string(img, config='--psm 6')
150
+ images.append(img)
151
+ except Exception as ocr_e:
152
+ logger.warning(f"OCR failed: {str(ocr_e)}")
153
+ if extracted:
154
+ lines = extracted.split("\n")
155
+ cleaned_lines = [line for line in lines if not re.match(r'^\s*(Page \d+|.*\d{4}-\d{4}|Copyright.*)\s*$', line, re.I)]
156
+ text += "\n".join(cleaned_lines) + "\n"
157
+ for char in page.chars:
158
+ if 'fontname' in char and 'mono' in char['fontname'].lower():
159
+ code_blocks.append(char['text'])
160
+ code_text = page.extract_text()
161
+ code_matches = re.finditer(r'(^\s{2,}.*?(?:\n\s{2,}.*?)*)', code_text, re.MULTILINE)
162
+ for match in code_matches:
163
+ code_blocks.append(match.group().strip())
164
+ tables = page.extract_tables()
165
+ if tables:
166
+ for table in tables:
167
+ text += "\n".join([" | ".join(map(str, row)) for row in table if row]) + "\n"
168
+ for obj in page.extract_words():
169
+ if obj.get('size', 0) > 12:
170
+ text += f"\n{obj['text']}\n"
171
+
172
+ code_text = "\n".join(code_blocks).strip()
173
+ if not text:
174
+ raise ValueError("No text extracted from PDF")
175
+
176
+ text_splitter = CharacterTextSplitter(separator="\n\n", chunk_size=250, chunk_overlap=40, keep_separator=True)
177
+ text_chunks = text_splitter.split_text(text)[:25]
178
+ code_chunks = text_splitter.split_text(code_text)[:10] if code_text else []
179
+
180
+ embeddings_model = load_embeddings_model()
181
+ if not embeddings_model:
182
+ return None, None, text, code_text, images
183
+
184
+ text_vector_store = FAISS.from_embeddings(
185
+ zip(text_chunks, [embeddings_model.encode(chunk, show_progress_bar=False, batch_size=128) for chunk in text_chunks]),
186
+ embeddings_model.encode
187
+ ) if text_chunks else None
188
+ code_vector_store = FAISS.from_embeddings(
189
+ zip(code_chunks, [embeddings_model.encode(chunk, show_progress_bar=False, batch_size=128) for chunk in code_chunks]),
190
+ embeddings_model.encode
191
+ ) if code_chunks else None
192
+
193
+ if text_vector_store:
194
+ text_vector_store = augment_vector_store(text_vector_store)
195
+
196
+ logger.info("PDF processed successfully")
197
+ return text_vector_store, code_vector_store, text, code_text, images
198
+ except Exception as e:
199
+ logger.error(f"PDF processing error: {str(e)}")
200
+ st.error(f"PDF error: {str(e)}")
201
+ return None, None, "", "", []
202
+
203
+ # Summarize PDF with ROUGE metrics and improved topic focus
204
+ def summarize_pdf(text):
205
+ logger.info("Generating summary")
206
+ try:
207
+ summary_pipeline = load_summary_pipeline()
208
+ if not summary_pipeline:
209
+ return "Summary model unavailable."
210
+
211
+ text_splitter = CharacterTextSplitter(separator="\n\n", chunk_size=250, chunk_overlap=40)
212
+ chunks = text_splitter.split_text(text)
213
+
214
+ # Hybrid search for relevant chunks
215
+ embeddings_model = load_embeddings_model()
216
+ if embeddings_model and chunks:
217
+ temp_vector_store = FAISS.from_embeddings(
218
+ zip(chunks, [embeddings_model.encode(chunk, show_progress_bar=False) for chunk in chunks]),
219
+ embeddings_model.encode
220
+ )
221
+ bm25 = BM25Okapi([chunk.split() for chunk in chunks])
222
+ query = "main topic and key points"
223
+ bm25_docs = bm25.get_top_n(query.split(), chunks, n=4)
224
+ faiss_docs = temp_vector_store.similarity_search(query, k=4)
225
+ selected_chunks = list(set(bm25_docs + [doc.page_content for doc in faiss_docs]))[:4]
226
+ else:
227
+ selected_chunks = chunks[:4]
228
+
229
+ summaries = []
230
+ for chunk in selected_chunks:
231
+ summary = summary_pipeline(f"Summarize the main topic and key points in detail: {chunk[:250]}", max_length=100, min_length=50, do_sample=False)[0]['summary_text']
232
+ summaries.append(summary.strip())
233
+
234
+ combined_summary = " ".join(summaries)
235
+ if len(combined_summary.split()) > 250:
236
+ combined_summary = " ".join(combined_summary.split()[:250])
237
+
238
+ word_count = len(combined_summary.split())
239
+ scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)
240
+ scores = scorer.score(text[:500], combined_summary)
241
+ logger.info(f"ROUGE scores: {scores}")
242
+
243
+ return f"**Main Topic Summary** ({word_count} words):\n{combined_summary}\n\n**ROUGE-1**: {scores['rouge1'].fmeasure:.2f}"
244
+ except Exception as e:
245
+ logger.error(f"Summary error: {str(e)}")
246
+ return f"Oops, something went wrong summarizing: {str(e)}"
247
+
248
+ # Answer question with hybrid search
249
+ def answer_question(text_vector_store, code_vector_store, query):
250
+ logger.info(f"Processing query: {query}")
251
+ try:
252
+ if not text_vector_store and not code_vector_store:
253
+ return "Please upload a PDF first!"
254
+
255
+ qa_pipeline = load_qa_pipeline()
256
+ if not qa_pipeline:
257
+ return "Sorry, the QA model is unavailable right now."
258
+
259
+ is_code_query = any(keyword in query.lower() for keyword in ["code", "script", "function", "programming", "give me code", "show code"])
260
+ if is_code_query and code_vector_store:
261
+ docs = code_vector_store.similarity_search(query, k=3)
262
+ code = "\n".join(doc.page_content for doc in docs)
263
+ explanation = qa_pipeline(f"Explain this code: {code[:500]}")[0]['generated_text']
264
+ return f"**Code**:\n```python\n{code}\n```\n**Explanation**:\n{explanation}"
265
+
266
+ vector_store = text_vector_store
267
+ if not vector_store:
268
+ return "No relevant content found for your query."
269
+
270
+ # Hybrid search: FAISS + BM25
271
+ text_chunks = [doc.page_content for doc in vector_store.similarity_search(query, k=10)]
272
+ bm25 = BM25Okapi([chunk.split() for chunk in text_chunks])
273
+ bm25_docs = bm25.get_top_n(query.split(), text_chunks, n=5)
274
+ faiss_docs = vector_store.similarity_search(query, k=5)
275
+ combined_docs = list(set(bm25_docs + [doc.page_content for doc in faiss_docs]))[:5]
276
+ context = "\n".join(combined_docs)
277
+
278
+ prompt = f"Use the following PDF content to answer the question accurately and concisely. Avoid speculation and focus on the provided context:\n\n{context}\n\nQuestion: {query}\nAnswer:"
279
+ response = qa_pipeline(prompt)[0]['generated_text']
280
+ logger.info("Answer generated")
281
+ return f"**Answer**:\n{response.strip()}\n\n**Source Context**:\n{context[:500]}..."
282
+ except Exception as e:
283
+ logger.error(f"Query error: {str(e)}")
284
+ return f"Sorry, something went wrong: {str(e)}"
285
+
286
+ # Streamlit UI
287
+ try:
288
+ st.set_page_config(page_title="Smart PDF Q&A", page_icon="📄", layout="wide")
289
+ st.markdown("""
290
+ <style>
291
+ .main { max-width: 900px; margin: 0 auto; padding: 20px; }
292
+ .sidebar { background-color: #f8f9fa; padding: 10px; border-radius: 5px; }
293
+ .message { margin: 10px 0; padding: 10px; border-radius: 5px; display: block; }
294
+ .user { background-color: #e6f3ff; }
295
+ .assistant { background-color: #f0f0f0; }
296
+ .dark .user { background-color: #2a2a72; color: #fff; }
297
+ .dark .assistant { background-color: #2e2e2e; color: #fff; }
298
+ .stButton>button { background-color: #4CAF50; color: white; border: none; padding: 8px 16px; border-radius: 5px; }
299
+ .stButton>button:hover { background-color: #45a049; }
300
+ pre { background-color: #f8f8f8; padding: 10px; border-radius: 5px; overflow-x: auto; }
301
+ .header { background: linear-gradient(90deg, #4CAF50, #81C784); color: white; padding: 10px; border-radius: 5px; text-align: center; }
302
+ .progress-bar { background-color: #e0e0e0; border-radius: 5px; height: 10px; }
303
+ .progress-fill { background-color: #4CAF50; height: 100%; border-radius: 5px; transition: width 0.5s ease; }
304
+ </style>
305
+ """, unsafe_allow_html=True)
306
+
307
+ st.markdown('<div class="header"><h1>Smart PDF Q&A</h1></div>', unsafe_allow_html=True)
308
+ st.markdown("Upload a PDF to ask questions, summarize (~150 words), or extract code with 'give me code'. Fast and friendly responses!")
309
+
310
+ # Initialize session state
311
+ if "messages" not in st.session_state:
312
+ st.session_state.messages = [{"role": "assistant", "content": "Hello! Upload a PDF and process it to start chatting."}]
313
+ if "text_vector_store" not in st.session_state:
314
+ st.session_state.text_vector_store = None
315
+ if "code_vector_store" not in st.session_state:
316
+ st.session_state.code_vector_store = None
317
+ if "pdf_text" not in st.session_state:
318
+ st.session_state.pdf_text = ""
319
+ if "code_text" not in st.session_state:
320
+ st.session_state.code_text = ""
321
+ if "images" not in st.session_state:
322
+ st.session_state.images = []
323
+
324
+ # Sidebar with toggle
325
+ with st.sidebar:
326
+ st.markdown('<div class="sidebar">', unsafe_allow_html=True)
327
+ theme = st.radio("Theme", ["Light", "Dark"], index=0)
328
+ dataset_name = st.selectbox("Select Dataset for Fine-Tuning", ["squad", "cnn_dailymail", "bigcode/the-stack"], index=0)
329
+ if st.button("Fine-Tune Model"):
330
+ progress_bar = st.progress(0)
331
+ for i in range(100):
332
+ time.sleep(0.008)
333
+ progress_bar.progress(i + 1)
334
+ dataset = load_and_prepare_dataset(dataset_name=dataset_name)
335
+ if dataset:
336
+ fine_tuned_pipeline = fine_tune_qa_model(dataset)
337
+ if fine_tuned_pipeline:
338
+ st.success("Model fine-tuned successfully!")
339
+ else:
340
+ st.error("Fine-tuning failed.")
341
+ if st.button("Clear Chat"):
342
  st.session_state.messages = []
343
+ st.experimental_rerun()
344
+ if st.button("Retry Summarization") and st.session_state.pdf_text:
345
+ progress_bar = st.progress(0)
346
+ with st.spinner("Retrying summarization..."):
347
+ for i in range(100):
348
+ time.sleep(0.008)
349
+ progress_bar.progress(i + 1)
350
+ summary = summarize_pdf(st.session_state.pdf_text)
351
+ st.session_state.messages.append({"role": "assistant", "content": summary})
352
+ st.markdown(summary, unsafe_allow_html=True)
353
+ st.markdown('</div>', unsafe_allow_html=True)
354
+
355
+ # PDF upload and processing
356
+ uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
357
+ col1, col2 = st.columns([1, 1])
358
+ with col1:
359
+ if st.button("Process PDF"):
360
+ progress_bar = st.progress(0)
361
+ with st.spinner("Processing PDF..."):
362
+ for i in range(100):
363
+ time.sleep(0.02)
364
+ progress_bar.progress(i + 1)
365
+ st.session_state.text_vector_store, st.session_state.code_vector_store, st.session_state.pdf_text, st.session_state.code_text, st.session_state.images = process_pdf(uploaded_file)
366
+ if st.session_state.text_vector_store or st.session_state.code_vector_store:
367
+ st.success("PDF processed! Ask away or summarize.")
368
+ st.session_state.messages = [{"role": "assistant", "content": "PDF processed! What would you like to know?"}]
369
+ else:
370
+ st.error("Failed to process PDF.")
371
+ with col2:
372
+ if st.button("Summarize PDF") and st.session_state.pdf_text:
373
+ progress_bar = st.progress(0)
374
+ with st.spinner("Summarizing..."):
375
+ for i in range(100):
376
+ time.sleep(0.008)
377
+ progress_bar.progress(i + 1)
378
+ summary = summarize_pdf(st.session_state.pdf_text)
379
+ st.session_state.messages.append({"role": "assistant", "content": summary})
380
+ st.markdown(summary, unsafe_allow_html=True)
381
 
382
+ # Chat interface
383
+ if st.session_state.text_vector_store or st.session_state.code_vector_store:
384
+ prompt = st.chat_input("Ask a question (e.g., 'Give me code' or 'What’s the main idea?'):")
385
+ if prompt:
386
+ st.session_state.messages.append({"role": "user", "content": prompt})
387
+ with st.chat_message("user"):
388
+ st.markdown(prompt)
389
+ with st.chat_message("assistant"):
390
+ progress_bar = st.progress(0)
391
+ with st.spinner('<div class="spinner">⏳ Processing...</div>'):
392
+ for i in range(100):
393
+ time.sleep(0.004)
394
+ progress_bar.progress(i + 1)
395
+ answer = answer_question(st.session_state.text_vector_store, st.session_state.code_vector_store, prompt)
396
+ st.markdown(answer, unsafe_allow_html=True)
397
+ st.session_state.messages.append({"role": "assistant", "content": answer})
398
+
399
+ # Display chat history
400
+ for message in st.session_state.messages:
401
+ with st.chat_message(message["role"]):
402
+ st.markdown(message["content"], unsafe_allow_html=True)
403
+
404
+ # Display extracted images
405
  if st.session_state.images:
406
+ st.header("Extracted Images")
407
+ for img in st.session_state.images:
408
+ st.image(img, caption="Extracted PDF Image", use_column_width=True)
409
+
410
+ # Download chat history
411
+ if st.session_state.messages:
412
+ chat_text = "\n".join(f"{m['role'].capitalize()}: {m['content']}" for m in st.session_state.messages)
413
+ st.download_button("Download Chat History", chat_text, "chat_history.txt")
414
+
415
+ except Exception as e:
416
+ logger.error(f"App initialization failed: {str(e)}")
417
+ st.error(f"App failed to start: {str(e)}. Check Spaces logs or contact support.")