Update app.py

app.py

@@ -24,7 +24,7 @@ logger = logging.getLogger(__name__)
 def load_embeddings_model():
     logger.info("Loading embeddings model")
     try:
-        return SentenceTransformer("all-MiniLM-L12-v2")
+        return SentenceTransformer("all-MiniLM-L6-v2")
     except Exception as e:
         logger.error(f"Embeddings load error: {str(e)}")
         st.error(f"Embedding model error: {str(e)}")
@@ -39,7 +39,7 @@ def load_qa_pipeline():
             fine_tuned_pipeline = fine_tune_qa_model(dataset)
             if fine_tuned_pipeline:
                 return fine_tuned_pipeline
-        return pipeline("text2text-generation", model="google/flan-t5-base", max_length=300)
+        return pipeline("text2text-generation", model="google/flan-t5-small", max_length=300)
     except Exception as e:
         logger.error(f"QA model load error: {str(e)}")
         st.error(f"QA model error: {str(e)}")
@@ -49,7 +49,7 @@ def load_qa_pipeline():
 def load_summary_pipeline():
     logger.info("Loading summary pipeline")
     try:
-        return pipeline("summarization", model="sshleifer/distilbart-cnn-6-6", max_length=150)
+        return pipeline("summarization", model="facebook/bart-large-cnn", max_length=250)
     except Exception as e:
         logger.error(f"Summary model load error: {str(e)}")
         st.error(f"Summary model error: {str(e)}")
@@ -79,7 +79,7 @@ def load_and_prepare_dataset(dataset_name="squad", max_samples=1000):
 def fine_tune_qa_model(dataset):
     logger.info("Starting fine-tuning")
     try:
-        model_name = "google/flan-t5-base"
+        model_name = "google/flan-t5-small"
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
         
@@ -93,13 +93,13 @@ def fine_tune_qa_model(dataset):
         
         training_args = TrainingArguments(
             output_dir="./fine_tuned_model",
-            num_train_epochs=2,
+            num_train_epochs = 2,
             per_device_train_batch_size=4,
             save_steps=500,
             logging_steps=100,
             evaluation_strategy="no",
             learning_rate=3e-5,
-            fp16=False,  # Set True if GPU available
+            fp16=False,
         )
         
         trainer = Trainer(
@@ -118,14 +118,14 @@ def fine_tune_qa_model(dataset):
         return None
 
 # Augment vector store with dataset
-def augment_vector_store(vector_store, dataset_name="squad", max_samples=500):
+def augment_vector_store(vector_store, dataset_name="squad", max_samples=300):
     logger.info(f"Augmenting vector store with dataset: {dataset_name}")
     try:
         dataset = load_dataset(dataset_name, split="train").select(range(min(max_samples, len(dataset))))
         chunks = [f"Context: {c}\nAnswer: {a['text'][0]}" for c, a in zip(dataset['context'], dataset['answers'])]
         embeddings_model = load_embeddings_model()
         if embeddings_model and vector_store:
-            embeddings = embeddings_model.encode(chunks, batch_size=32, show_progress_bar=False)
+            embeddings = embeddings_model.encode(chunks, batch_size=128, show_progress_bar=False)
             vector_store.add_embeddings(zip(chunks, embeddings))
         return vector_store
     except Exception as e:
@@ -139,16 +139,19 @@ def process_pdf(uploaded_file):
         text = ""
         code_blocks = []
         with pdfplumber.open(BytesIO(uploaded_file.getvalue())) as pdf:
-            for page in pdf.pages[:20]:
+            for page in pdf.pages[:8]:
                 extracted = page.extract_text(layout=False)
-                if not extracted:  # OCR fallback for scanned PDFs
+                if not extracted:
                     try:
                         img = page.to_image(resolution=150).original
                         extracted = pytesseract.image_to_string(img, config='--psm 6')
                     except Exception as ocr_e:
                         logger.warning(f"OCR failed: {str(ocr_e)}")
                 if extracted:
-                    text += extracted + "\n"
+                    # Clean text: remove headers/footers (simple heuristic)
+                    lines = extracted.split("\n")
+                    cleaned_lines = [line for line in lines if not re.match(r'^\s*(Page \d+|.*\d{4}-\d{4}|Copyright.*)\s*$', line, re.I)]
+                    text += "\n".join(cleaned_lines) + "\n"
                 for char in page.chars:
                     if 'fontname' in char and 'mono' in char['fontname'].lower():
                         code_blocks.append(char['text'])
@@ -168,20 +171,20 @@ def process_pdf(uploaded_file):
         if not text:
             raise ValueError("No text extracted from PDF")
         
-        text_splitter = CharacterTextSplitter(separator="\n\n", chunk_size=400, chunk_overlap=80, keep_separator=True)
-        text_chunks = text_splitter.split_text(text)[:80]
-        code_chunks = text_splitter.split_text(code_text)[:40] if code_text else []
+        text_splitter = CharacterTextSplitter(separator="\n\n", chunk_size=250, chunk_overlap=40, keep_separator=True)
+        text_chunks = text_splitter.split_text(text)[:25]
+        code_chunks = text_splitter.split_text(code_text)[:10] if code_text else []
         
         embeddings_model = load_embeddings_model()
         if not embeddings_model:
             return None, None, text, code_text
         
         text_vector_store = FAISS.from_embeddings(
-            zip(text_chunks, [embeddings_model.encode(chunk, show_progress_bar=False) for chunk in text_chunks]),
+            zip(text_chunks, [embeddings_model.encode(chunk, show_progress_bar=False, batch_size=128) for chunk in text_chunks]),
             embeddings_model.encode
         ) if text_chunks else None
         code_vector_store = FAISS.from_embeddings(
-            zip(code_chunks, [embeddings_model.encode(chunk, show_progress_bar=False) for chunk in code_chunks]),
+            zip(code_chunks, [embeddings_model.encode(chunk, show_progress_bar=False, batch_size=128) for chunk in code_chunks]),
             embeddings_model.encode
         ) if code_chunks else None
         
@@ -195,7 +198,7 @@ def process_pdf(uploaded_file):
         st.error(f"PDF error: {str(e)}")
         return None, None, "", ""
 
-# Summarize PDF with ROUGE metrics
+# Summarize PDF with ROUGE metrics and improved topic focus
 def summarize_pdf(text):
     logger.info("Generating summary")
     try:
@@ -203,23 +206,39 @@ def summarize_pdf(text):
         if not summary_pipeline:
             return "Summary model unavailable."
         
-        text_splitter = CharacterTextSplitter(separator="\n\n", chunk_size=400, chunk_overlap=50)
-        chunks = text_splitter.split_text(text)[:2]
-        summaries = []
+        text_splitter = CharacterTextSplitter(separator="\n\n", chunk_size=250, chunk_overlap=40)
+        chunks = text_splitter.split_text(text)
+        
+        # Hybrid search for relevant chunks
+        embeddings_model = load_embeddings_model()
+        if embeddings_model and chunks:
+            temp_vector_store = FAISS.from_embeddings(
+                zip(chunks, [embeddings_model.encode(chunk, show_progress_bar=False) for chunk in chunks]),
+                embeddings_model.encode
+            )
+            bm25 = BM25Okapi([chunk.split() for chunk in chunks])
+            query = "main topic and key points"
+            bm25_docs = bm25.get_top_n(query.split(), chunks, n=4)
+            faiss_docs = temp_vector_store.similarity_search(query, k=4)
+            selected_chunks = list(set(bm25_docs + [doc.page_content for doc in faiss_docs]))[:4]
+        else:
+            selected_chunks = chunks[:4]
         
-        for chunk in chunks:
-            summary = summary_pipeline(chunk[:400], max_length=100, min_length=30, do_sample=False)[0]['summary_text']
+        summaries = []
+        for chunk in selected_chunks:
+            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']
             summaries.append(summary.strip())
         
         combined_summary = " ".join(summaries)
-        if len(combined_summary.split()) > 150:
-            combined_summary = " ".join(combined_summary.split()[:150])
+        if len(combined_summary.split()) > 250:
+            combined_summary = " ".join(combined_summary.split()[:250])
         
+        word_count = len(combined_summary.split())
         scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)
-        scores = scorer.score(text[:400], combined_summary)
+        scores = scorer.score(text[:500], combined_summary)
         logger.info(f"ROUGE scores: {scores}")
         
-        return f"**Summary**:\n{combined_summary}\n\n**ROUGE-1**: {scores['rouge1'].fmeasure:.2f}"
+        return f"**Main Topic Summary** ({word_count} words):\n{combined_summary}\n\n**ROUGE-1**: {scores['rouge1'].fmeasure:.2f}"
     except Exception as e:
         logger.error(f"Summary error: {str(e)}")
         return f"Oops, something went wrong summarizing: {str(e)}"
@@ -285,7 +304,7 @@ try:
     """, unsafe_allow_html=True)
 
     st.markdown('<div class="header"><h1>Smart PDF Q&A</h1></div>', unsafe_allow_html=True)
-    st.markdown("Upload a PDF to ask questions, summarize (~150 words), or extract code with 'give me code'. Fast, accurate, and smooth!")
+    st.markdown("Upload a PDF to ask questions, get a ~200-word summary, or extract code with 'give me code'. Optimized for speed and accuracy!")
 
     # Initialize session state
     if "messages" not in st.session_state:
@@ -307,7 +326,7 @@ try:
         if st.button("Fine-Tune Model"):
             progress_bar = st.progress(0)
             for i in range(100):
-                time.sleep(0.02)
+                time.sleep(0.008)
                 progress_bar.progress(i + 1)
             dataset = load_and_prepare_dataset(dataset_name=dataset_name)
             if dataset:
@@ -319,6 +338,15 @@ try:
         if st.button("Clear Chat"):
             st.session_state.messages = []
             st.experimental_rerun()
+        if st.button("Retry Summarization") and st.session_state.pdf_text:
+            progress_bar = st.progress(0)
+            with st.spinner("Retrying summarization..."):
+                for i in range(100):
+                    time.sleep(0.008)
+                    progress_bar.progress(i + 1)
+                summary = summarize_pdf(st.session_state.pdf_text)
+                st.session_state.messages.append({"role": "assistant", "content": summary})
+                st.markdown(summary, unsafe_allow_html=True)
         st.markdown('</div>', unsafe_allow_html=True)
 
     # PDF upload and processing
@@ -329,7 +357,7 @@ try:
             progress_bar = st.progress(0)
             with st.spinner("Processing PDF..."):
                 for i in range(100):
-                    time.sleep(0.05)
+                    time.sleep(0.02)
                     progress_bar.progress(i + 1)
                 st.session_state.text_vector_store, st.session_state.code_vector_store, st.session_state.pdf_text, st.session_state.code_text = process_pdf(uploaded_file)
                 if st.session_state.text_vector_store or st.session_state.code_vector_store:
@@ -342,7 +370,7 @@ try:
             progress_bar = st.progress(0)
             with st.spinner("Summarizing..."):
                 for i in range(100):
-                    time.sleep(0.02)
+                    time.sleep(0.008)
                     progress_bar.progress(i + 1)
                 summary = summarize_pdf(st.session_state.pdf_text)
                 st.session_state.messages.append({"role": "assistant", "content": summary})
@@ -360,7 +388,7 @@ try:
                 progress_bar = st.progress(0)
                 with st.spinner('<div class="spinner">⏳ Processing...</div>'):
                     for i in range(100):
-                        time.sleep(0.01)
+                        time.sleep(0.004)
                         progress_bar.progress(i + 1)
                     answer = answer_question(st.session_state.text_vector_store, st.session_state.code_vector_store, prompt)
                 st.markdown(answer, unsafe_allow_html=True)