# -*- coding: utf-8 -*- """ @title: PDF Chat Assistant @author: Your Name """ import streamlit as st import fitz import torch from sentence_transformers import SentenceTransformer, util from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer from tqdm.auto import tqdm import textwrap import re from spacy.lang.en import English from threading import Thread import os # Configuration MODEL_NAME = "all-mpnet-base-v2" LLM_MODEL_ID = "google/gemma-2b-it" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" NUM_RESULTS = 5 MIN_TOKEN_LENGTH = 30 # Set cache directories os.environ['TRANSFORMERS_CACHE'] = '/app/.cache' os.environ['HF_HOME'] = '/app/.cache' # Initialize session state if 'processed' not in st.session_state: st.session_state.processed = False if 'embeddings' not in st.session_state: st.session_state.embeddings = None if 'pages_and_chunks' not in st.session_state: st.session_state.pages_and_chunks = [] if 'messages' not in st.session_state: st.session_state.messages = [] # Helper functions def text_formatter(text: str) -> str: cleaned_text = text.replace("\n", " ").strip() cleaned_text = re.sub(r'\s+', ' ', cleaned_text) return cleaned_text def split_list(input_list: list, slice_size: int = 10) -> list[list[str]]: return [input_list[i:i + slice_size] for i in range(0, len(input_list), slice_size)] # PDF Processing def process_pdf(uploaded_file): try: doc = fitz.open(stream=uploaded_file.read(), filetype="pdf") pages_and_texts = [] with st.spinner("📄 Processing PDF..."): for page_number, page in tqdm(enumerate(doc)): text = page.get_text() text = text_formatter(text) pages_and_texts.append({ "page_number": page_number, "text": text }) with st.spinner("🔪 Chunking text..."): nlp = English() nlp.add_pipe("sentencizer") for item in tqdm(pages_and_texts): item['sentences'] = [str(s) for s in nlp(item["text"]).sents] item["sentence_chunks"] = split_list(item["sentences"]) pages_and_chunks = [] for item in tqdm(pages_and_texts): for sentence_chunk in item["sentence_chunks"]: chunk_text = " ".join(sentence_chunk).replace(" ", " ").strip() pages_and_chunks.append({ "page_number": item["page_number"], "sentence_chunk": chunk_text, "chunk_token_count": len(chunk_text)/4 }) return [c for c in pages_and_chunks if c["chunk_token_count"] > MIN_TOKEN_LENGTH] except Exception as e: st.error(f"❌ Error processing PDF: {str(e)}") return None # Model Loading @st.cache_resource def load_models(): try: # Load embedding model embedding_model = SentenceTransformer(MODEL_NAME, device=DEVICE) # Load LLM quantization_config = BitsAndBytesConfig(load_in_4bit=True) tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID) llm_model = AutoModelForCausalLM.from_pretrained( LLM_MODEL_ID, quantization_config=quantization_config, device_map="auto" ) return embedding_model, tokenizer, llm_model except Exception as e: st.error(f"❌ Error loading models: {str(e)}") return None, None, None # Chat Functions def format_prompt(query, context_items): system_prompt = """You are an AI assistant that answers questions based on PDF documents. Use only the provided context to answer questions. Be technical and detailed in your responses.""" context = "\n".join([f"Context {i+1}: {c['sentence_chunk']}" for i, c in enumerate(context_items)]) return f"""user {system_prompt} Question: {query} Document Context: {context} Please provide a detailed answer using the document context. model """ def rag_response(query): try: embedding_model, tokenizer, llm_model = load_models() if None in (embedding_model, tokenizer, llm_model): return "Error: Models not loaded properly" # Retrieve relevant context query_embedding = embedding_model.encode(query, convert_to_tensor=True) scores = util.dot_score(query_embedding, st.session_state.embeddings)[0] top_results = torch.topk(scores, k=NUM_RESULTS) context_items = [st.session_state.pages_and_chunks[i] for i in top_results[1]] # Create streamer streamer = TextIteratorStreamer(tokenizer, skip_prompt=True) # Format prompt prompt = format_prompt(query, context_items) # Generation parameters inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE) generation_kwargs = dict( inputs, streamer=streamer, max_new_tokens=1000, temperature=0.7, do_sample=True ) # Start generation in separate thread thread = Thread(target=llm_model.generate, kwargs=generation_kwargs) thread.start() # Stream the response full_response = "" with st.empty(): for new_text in streamer: full_response += new_text st.markdown(full_response + "▌") st.markdown(full_response) return full_response except Exception as e: return f"Error generating response: {str(e)}" # Streamlit UI st.title("📚 PDF Chat Assistant") st.caption("Chat with your documents using AI") # Sidebar for PDF Upload with st.sidebar: st.header("Document Setup") uploaded_file = st.file_uploader("Upload PDF", type=["pdf"]) if uploaded_file and not st.session_state.processed: st.session_state.messages = [] st.session_state.pages_and_chunks = process_pdf(uploaded_file) if st.session_state.pages_and_chunks: embedding_model, _, _ = load_models() if embedding_model: with st.spinner("🔧 Generating embeddings..."): texts = [c["sentence_chunk"] for c in st.session_state.pages_and_chunks] st.session_state.embeddings = torch.tensor( embedding_model.encode(texts, convert_to_tensor=True), device=DEVICE ) st.session_state.processed = True # Chat Interface if st.session_state.processed: # Display chat messages for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # Handle user input if prompt := st.chat_input("Ask about your document..."): # Add user message to chat history st.session_state.messages.append({"role": "user", "content": prompt}) # Display user message with st.chat_message("user"): st.markdown(prompt) # Generate and display AI response with st.chat_message("assistant"): response = rag_response(prompt) st.session_state.messages.append({"role": "assistant", "content": response}) else: st.info("👉 Please upload a PDF document to get started") # Add footer st.sidebar.markdown("---") st.sidebar.markdown("Powered by [Gemma](https://ai.google.dev/gemma) • [Hugging Face](https://huggingface.co)")