import streamlit as st
import logging
import os
from io import BytesIO
import pdfplumber
from pdf2image import convert_from_bytes
from PIL import Image
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.vectorstores import FAISS
from sentence_transformers import SentenceTransformer
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments
from datasets import load_dataset
from rank_bm25 import BM25Okapi
from rouge_score import rouge_scorer
import re
import time
import pytesseract
# Setup logging for Spaces
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Lazy load models
@st.cache_resource(ttl=1800)
def load_embeddings_model():
logger.info("Loading embeddings model")
try:
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)}")
return None
@st.cache_resource(ttl=1800)
def load_qa_pipeline():
logger.info("Loading QA pipeline")
try:
dataset = load_and_prepare_dataset()
if dataset:
fine_tuned_pipeline = fine_tune_qa_model(dataset)
if fine_tuned_pipeline:
return fine_tuned_pipeline
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)}")
return None
@st.cache_resource(ttl=1800)
def load_summary_pipeline():
logger.info("Loading summary pipeline")
try:
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)}")
return None
# Load and prepare dataset (e.g., SQuAD)
@st.cache_data(ttl=3600)
def load_and_prepare_dataset(dataset_name="squad", max_samples=1000):
logger.info(f"Loading dataset: {dataset_name}")
try:
dataset = load_dataset(dataset_name, split="train[:80%]")
dataset = dataset.shuffle(seed=42).select(range(min(max_samples, len(dataset))))
def preprocess(examples):
inputs = [f"question: {q} context: {c}" for q, c in zip(examples['question'], examples['context'])]
targets = examples['answers']['text']
return {'input_text': inputs, 'target_text': [t[0] if t else "" for t in targets]}
dataset = dataset.map(preprocess, batched=True, remove_columns=dataset.column_names)
return dataset
except Exception as e:
logger.error(f"Dataset load error: {str(e)}")
return None
# Fine-tune QA model
@st.cache_resource(ttl=3600)
def fine_tune_qa_model(dataset):
logger.info("Starting fine-tuning")
try:
model_name = "google/flan-t5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
def tokenize_function(examples):
model_inputs = tokenizer(examples['input_text'], max_length=512, truncation=True, padding="max_length")
labels = tokenizer(examples['target_text'], max_length=128, truncation=True, padding="max_length")
model_inputs["labels"] = labels["input_ids"]
return model_inputs
tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=['input_text', 'target_text'])
training_args = TrainingArguments(
output_dir="./fine_tuned_model",
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,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
)
trainer.train()
model.save_pretrained("./fine_tuned_model")
tokenizer.save_pretrained("./fine_tuned_model")
logger.info("Fine-tuning complete")
return pipeline("text2text-generation", model="./fine_tuned_model", tokenizer="./fine_tuned_model", max_length=300)
except Exception as e:
logger.error(f"Fine-tuning error: {str(e)}")
return None
# Augment vector store with dataset
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=128, show_progress_bar=False)
vector_store.add_embeddings(zip(chunks, embeddings))
return vector_store
except Exception as e:
logger.error(f"Vector store augmentation error: {str(e)}")
return vector_store
# Process PDF with enhanced extraction and OCR fallback
def process_pdf(uploaded_file):
logger.info("Processing PDF with enhanced extraction")
try:
text = ""
code_blocks = []
images = []
with pdfplumber.open(BytesIO(uploaded_file.getvalue())) as pdf:
for page in pdf.pages[:8]:
extracted = page.extract_text(layout=False)
if not extracted:
try:
img = page.to_image(resolution=150).original
extracted = pytesseract.image_to_string(img, config='--psm 6')
images.append(img)
except Exception as ocr_e:
logger.warning(f"OCR failed: {str(ocr_e)}")
if extracted:
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'])
code_text = page.extract_text()
code_matches = re.finditer(r'(^\s{2,}.*?(?:\n\s{2,}.*?)*)', code_text, re.MULTILINE)
for match in code_matches:
code_blocks.append(match.group().strip())
tables = page.extract_tables()
if tables:
for table in tables:
text += "\n".join([" | ".join(map(str, row)) for row in table if row]) + "\n"
for obj in page.extract_words():
if obj.get('size', 0) > 12:
text += f"\n{obj['text']}\n"
code_text = "\n".join(code_blocks).strip()
if not text:
raise ValueError("No text extracted from PDF")
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, images
text_vector_store = FAISS.from_embeddings(
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, batch_size=128) for chunk in code_chunks]),
embeddings_model.encode
) if code_chunks else None
if text_vector_store:
text_vector_store = augment_vector_store(text_vector_store)
logger.info("PDF processed successfully")
return text_vector_store, code_vector_store, text, code_text, images
except Exception as e:
logger.error(f"PDF processing error: {str(e)}")
st.error(f"PDF error: {str(e)}")
return None, None, "", "", []
# Summarize PDF with ROUGE metrics and improved topic focus
def summarize_pdf(text):
logger.info("Generating summary")
try:
summary_pipeline = load_summary_pipeline()
if not summary_pipeline:
return "Summary model unavailable."
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]
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()) > 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[:500], combined_summary)
logger.info(f"ROUGE scores: {scores}")
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)}"
# Answer question with hybrid search
def answer_question(text_vector_store, code_vector_store, query):
logger.info(f"Processing query: {query}")
try:
if not text_vector_store and not code_vector_store:
return "Please upload a PDF first!"
qa_pipeline = load_qa_pipeline()
if not qa_pipeline:
return "Sorry, the QA model is unavailable right now."
is_code_query = any(keyword in query.lower() for keyword in ["code", "script", "function", "programming", "give me code", "show code"])
if is_code_query and code_vector_store:
docs = code_vector_store.similarity_search(query, k=3)
code = "\n".join(doc.page_content for doc in docs)
explanation = qa_pipeline(f"Explain this code: {code[:500]}")[0]['generated_text']
return f"**Code**:\n```python\n{code}\n```\n**Explanation**:\n{explanation}"
vector_store = text_vector_store
if not vector_store:
return "No relevant content found for your query."
# Hybrid search: FAISS + BM25
text_chunks = [doc.page_content for doc in vector_store.similarity_search(query, k=10)]
bm25 = BM25Okapi([chunk.split() for chunk in text_chunks])
bm25_docs = bm25.get_top_n(query.split(), text_chunks, n=5)
faiss_docs = vector_store.similarity_search(query, k=5)
combined_docs = list(set(bm25_docs + [doc.page_content for doc in faiss_docs]))[:5]
context = "\n".join(combined_docs)
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:"
response = qa_pipeline(prompt)[0]['generated_text']
logger.info("Answer generated")
return f"**Answer**:\n{response.strip()}\n\n**Source Context**:\n{context[:500]}..."
except Exception as e:
logger.error(f"Query error: {str(e)}")
return f"Sorry, something went wrong: {str(e)}"
# Streamlit UI
try:
st.set_page_config(page_title="Smart PDF Q&A", page_icon="đź“„", layout="wide")
st.markdown("""
""", unsafe_allow_html=True)
st.markdown('', unsafe_allow_html=True)
st.markdown("Upload a PDF to ask questions, summarize (~150 words), or extract code with 'give me code'. Fast and friendly responses!")
# Initialize session state
if "messages" not in st.session_state:
st.session_state.messages = [{"role": "assistant", "content": "Hello! Upload a PDF and process it to start chatting."}]
if "text_vector_store" not in st.session_state:
st.session_state.text_vector_store = None
if "code_vector_store" not in st.session_state:
st.session_state.code_vector_store = None
if "pdf_text" not in st.session_state:
st.session_state.pdf_text = ""
if "code_text" not in st.session_state:
st.session_state.code_text = ""
if "images" not in st.session_state:
st.session_state.images = []
# Sidebar with toggle
with st.sidebar:
st.markdown('', unsafe_allow_html=True)
# PDF upload and processing
uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
col1, col2 = st.columns([1, 1])
with col1:
if st.button("Process PDF"):
progress_bar = st.progress(0)
with st.spinner("Processing PDF..."):
for i in range(100):
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, st.session_state.images = process_pdf(uploaded_file)
if st.session_state.text_vector_store or st.session_state.code_vector_store:
st.success("PDF processed! Ask away or summarize.")
st.session_state.messages = [{"role": "assistant", "content": "PDF processed! What would you like to know?"}]
else:
st.error("Failed to process PDF.")
with col2:
if st.button("Summarize PDF") and st.session_state.pdf_text:
progress_bar = st.progress(0)
with st.spinner("Summarizing..."):
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)
# Chat interface
if st.session_state.text_vector_store or st.session_state.code_vector_store:
prompt = st.chat_input("Ask a question (e.g., 'Give me code' or 'What’s the main idea?'):")
if prompt:
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
progress_bar = st.progress(0)
with st.spinner('⏳ Processing...
'):
for i in range(100):
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)
st.session_state.messages.append({"role": "assistant", "content": answer})
# Display chat history
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"], unsafe_allow_html=True)
# Display extracted images
if st.session_state.images:
st.header("Extracted Images")
for img in st.session_state.images:
st.image(img, caption="Extracted PDF Image", use_column_width=True)
# Download chat history
if st.session_state.messages:
chat_text = "\n".join(f"{m['role'].capitalize()}: {m['content']}" for m in st.session_state.messages)
st.download_button("Download Chat History", chat_text, "chat_history.txt")
except Exception as e:
logger.error(f"App initialization failed: {str(e)}")
st.error(f"App failed to start: {str(e)}. Check Spaces logs or contact support.")