src/streamlit_app.py

# -*- 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"""<start_of_turn>user
{system_prompt}

Question: {query}
Document Context: {context}

Please provide a detailed answer using the document context.<end_of_turn>
<start_of_turn>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)")