sunbal7's picture
Update app.py
64794fb verified
raw
history blame
11.6 kB
import streamlit as st
import torch
import numpy as np
import faiss
import time
import re
from typing import List, Tuple
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from sentence_transformers import SentenceTransformer
import fitz # PyMuPDF
import docx2txt
from langchain_text_splitters import RecursiveCharacterTextSplitter
from io import BytesIO
# ------------------------
# Configuration
# ------------------------
MODEL_NAME = "ibm-granite/granite-3.1-1b-a400m-instruct"
EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
CHUNK_SIZE = 1024 # Increased for better context
CHUNK_OVERLAP = 128
MAX_FILE_SIZE_MB = 10
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ------------------------
# Model Loading with Quantization
# ------------------------
@st.cache_resource
def load_models():
try:
# Configure quantization for CPU deployment
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
) if DEVICE == "cpu" else None
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
revision="main"
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
revision="main",
device_map="auto",
quantization_config=quant_config,
torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
low_cpu_mem_usage=True
).eval()
# Load embedding model with FP16 optimization
embedder = SentenceTransformer(
EMBED_MODEL,
device=DEVICE,
device_kwargs={"keep_all_models": True}
)
if DEVICE == "cuda":
embedder = embedder.half()
return tokenizer, model, embedder
except Exception as e:
st.error(f"Model loading failed: {str(e)}")
st.stop()
# ------------------------
# Enhanced Text Processing
# ------------------------
def clean_text(text: str) -> str:
"""Advanced text cleaning with multiple normalization steps"""
text = re.sub(r'\s+', ' ', text) # Remove extra whitespace
text = re.sub(r'[^\x00-\x7F]+', ' ', text) # Remove non-ASCII
text = re.sub(r'\bPage \d+\b', '', text) # Remove page numbers
text = re.sub(r'http\S+', '', text) # Remove URLs
text = re.sub(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '', text) # Remove emails
return text.strip()
def extract_text(file: BytesIO) -> Tuple[str, List[str]]:
"""Improved text extraction with format-specific handling"""
try:
if file.size > MAX_FILE_SIZE_MB * 1024 * 1024:
raise ValueError(f"File size exceeds {MAX_FILE_SIZE_MB}MB limit")
file_type = file.type
text = ""
if file_type == "application/pdf":
doc = fitz.open(stream=file.read(), filetype="pdf")
text = "\n".join([page.get_text("text", flags=fitz.TEXT_PRESERVE_WHITESPACE) for page in doc])
# Extract images metadata for future multimodal expansion
images = [img for page in doc for img in page.get_images()]
if images:
st.session_state.images = images
elif file_type == "text/plain":
text = file.read().decode("utf-8", errors="replace")
elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
text = docx2txt.process(file)
else:
raise ValueError("Unsupported file type")
return clean_text(text)
except Exception as e:
st.error(f"Text extraction failed: {str(e)}")
st.stop()
def semantic_chunking(text: str) -> List[str]:
"""Context-aware text splitting with metadata tracking"""
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
length_function=len,
add_start_index=True
)
chunks = splitter.split_text(text)
return chunks
# ------------------------
# Enhanced Vector Indexing
# ------------------------
def build_faiss_index(chunks: List[str], embedder) -> faiss.Index:
"""Build optimized FAISS index with error handling"""
try:
embeddings = embedder.encode(
chunks,
batch_size=32,
show_progress_bar=True,
convert_to_tensor=True
)
if DEVICE == "cuda":
embeddings = embeddings.cpu().numpy()
else:
embeddings = embeddings.numpy()
dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
faiss.normalize_L2(embeddings)
index.add(embeddings)
return index
except Exception as e:
st.error(f"Index creation failed: {str(e)}")
st.stop()
# ------------------------
# Improved Generation Functions
# ------------------------
def format_prompt(system_prompt: str, user_input: str) -> str:
"""Structured prompt formatting for better model performance"""
return f"""<|system|>
{system_prompt}
<|user|>
{user_input}
<|assistant|>
"""
def generate_summary(text: str, tokenizer, model) -> str:
"""Hierarchical summarization with chunk processing"""
try:
# First-stage summary
chunks = [text[i:i+3000] for i in range(0, len(text), 3000)]
summaries = []
for chunk in chunks:
prompt = format_prompt(
"Generate a detailed summary of this text excerpt:",
chunk[:2500]
)
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
outputs = model.generate(
**inputs,
max_new_tokens=300,
temperature=0.3,
do_sample=True
)
summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
summaries.append(summary.split("<|assistant|>")[-1].strip())
# Final synthesis
final_prompt = format_prompt(
"Synthesize these summaries into a comprehensive overview:",
"\n".join(summaries)
)
inputs = tokenizer(final_prompt, return_tensors="pt").to(DEVICE)
outputs = model.generate(
**inputs,
max_new_tokens=500,
temperature=0.4,
do_sample=True
)
return tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
except Exception as e:
st.error(f"Summarization failed: {str(e)}")
return "Summary generation failed"
def retrieve_context(query: str, index, chunks: List[str], embedder, top_k: int = 3) -> str:
"""Enhanced retrieval with score thresholding"""
query_embed = embedder.encode([query], convert_to_tensor=True)
if DEVICE == "cuda":
query_embed = query_embed.cpu().numpy()
else:
query_embed = query_embed.numpy()
faiss.normalize_L2(query_embed)
scores, indices = index.search(query_embed, top_k*2) # Retrieve extra for filtering
# Apply similarity threshold
valid_indices = [i for i, score in zip(indices[0], scores[0]) if score > 0.35]
return " ".join([chunks[i] for i in valid_indices[:top_k]])
# ------------------------
# Streamlit UI Improvements
# ------------------------
def main():
st.set_page_config(
page_title="RAG Book Analyzer Pro",
layout="wide",
initial_sidebar_state="expanded"
)
# Initialize session state
if "processed" not in st.session_state:
st.session_state.processed = False
if "index" not in st.session_state:
st.session_state.index = None
# Load models once
tokenizer, model, embedder = load_models()
# Sidebar controls
with st.sidebar:
st.header("Settings")
top_k = st.slider("Number of context passages", 1, 5, 3)
temp = st.slider("Generation Temperature", 0.1, 1.0, 0.4)
# Main interface
st.title("πŸ“š Advanced Book Analyzer")
st.write("Upload technical manuals, research papers, or books for deep analysis")
uploaded_file = st.file_uploader(
"Choose a document",
type=["pdf", "txt", "docx"],
accept_multiple_files=False
)
if uploaded_file and not st.session_state.processed:
with st.spinner("Analyzing document..."):
start_time = time.time()
# Process document
text = extract_text(uploaded_file)
chunks = semantic_chunking(text)
index = build_faiss_index(chunks, embedder)
# Store in session state
st.session_state.update({
"chunks": chunks,
"index": index,
"processed": True,
"text": text
})
st.success(f"Processed {len(chunks)} chunks in {time.time()-start_time:.1f}s")
if st.session_state.processed:
# Summary section
with st.expander("Document Summary", expanded=True):
summary = generate_summary(st.session_state.text, tokenizer, model)
st.markdown(summary)
# Q&A Section
st.divider()
col1, col2 = st.columns([3, 1])
with col1:
query = st.text_input("Ask about the document:", placeholder="What are the key findings...")
with col2:
show_context = st.checkbox("Show context sources")
if query:
with st.spinner("Searching document..."):
context = retrieve_context(
query,
st.session_state.index,
st.session_state.chunks,
embedder,
top_k=top_k
)
if not context:
st.warning("No relevant context found in document")
return
with st.expander("Generated Answer", expanded=True):
answer = generate_answer(query, context, tokenizer, model, temp)
st.markdown(answer)
if show_context:
st.divider()
st.subheader("Source Context")
st.write(context)
def generate_answer(query: str, context: str, tokenizer, model, temp: float) -> str:
"""Improved answer generation with context validation"""
try:
prompt = format_prompt(
f"""Answer the question using only the provided context.
Follow these rules:
1. Be precise and factual
2. If unsure, say 'The document does not specify'
3. Use bullet points when listing items
4. Keep answers under 3 sentences
Context: {context[:2000]}""",
query
)
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
outputs = model.generate(
**inputs,
max_new_tokens=400,
temperature=temp,
top_p=0.9,
repetition_penalty=1.2,
do_sample=True
)
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
return answer.split("<|assistant|>")[-1].strip()
except Exception as e:
st.error(f"Generation failed: {str(e)}")
return "Unable to generate answer"
if __name__ == "__main__":
main()