Update app.py

app.py

@@ -4,24 +4,25 @@ import numpy as np
 import faiss
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
-import fitz  # PyMuPDF for better PDF extraction
+import fitz  # PyMuPDF for PDF extraction
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-
 # Configuration
 MODEL_NAME = "ibm-granite/granite-3.1-1b-a400m-instruct"
+EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
+CHUNK_SIZE = 512
+CHUNK_OVERLAP = 64
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 @st.cache_resource
-def load_model():
+def load_models():
     try:
-        # Load with explicit configuration
+        # Load tokenizer and generative model with trust_remote_code enabled
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
             trust_remote_code=True,
             revision="main"
         )
-        
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             device_map="auto" if DEVICE == "cuda" else None,
@@ -29,26 +30,9 @@ def load_model():
             trust_remote_code=True,
             revision="main",
             low_cpu_mem_usage=True
-        )
-        return model, tokenizer
-    except Exception as e:
-        st.error(f"Model loading failed: {str(e)}")
-        st.stop()
-
-model, tokenizer = load_model()
-
-
-def load_models():
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            device_map="auto" if DEVICE == "cuda" else None,
-            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-            low_cpu_mem_usage=True
         ).eval()
         
+        # Load embedding model for FAISS
         embedder = SentenceTransformer(EMBED_MODEL, device=DEVICE)
         return tokenizer, model, embedder
     
@@ -58,7 +42,7 @@ def load_models():
 
 tokenizer, model, embedder = load_models()
 
-# Improved text processing
+# Improved text processing: splits text into chunks
 def process_text(text):
     splitter = RecursiveCharacterTextSplitter(
         chunk_size=CHUNK_SIZE,
@@ -67,7 +51,7 @@ def process_text(text):
     )
     return splitter.split_text(text)
 
-# Enhanced PDF extraction
+# Enhanced PDF extraction using PyMuPDF
 def extract_pdf_text(uploaded_file):
     try:
         doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
@@ -83,28 +67,29 @@ def generate_summary(text):
     
     for chunk in chunks:
         prompt = f"""<|user|>
-        Summarize this text section focusing on key themes, characters, and plot points:
-        {chunk[:2000]}
-        <|assistant|>
-        """
-        
+Summarize this text section focusing on key themes, characters, and plot points:
+{chunk[:2000]}
+<|assistant|>
+"""
         inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
         outputs = model.generate(**inputs, max_new_tokens=300, temperature=0.3)
-        summaries.append(tokenizer.decode(outputs[0], skip_special_tokens=True))
+        summary_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        summaries.append(summary_text)
     
-    # Combine summaries
+    # Combine individual summaries into one comprehensive summary
     combined = "\n".join(summaries)
     final_prompt = f"""<|user|>
-    Combine these section summaries into a coherent book summary:
-    {combined}
-    <|assistant|>
-    The comprehensive summary is:"""
+Combine these section summaries into a coherent book summary:
+{combined}
+<|assistant|>
+The comprehensive summary is:"""
     
     inputs = tokenizer(final_prompt, return_tensors="pt").to(DEVICE)
     outputs = model.generate(**inputs, max_new_tokens=500, temperature=0.5)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).split(":")[-1].strip()
+    full_summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return full_summary.split(":")[-1].strip()
 
-# Enhanced retrieval system
+# Enhanced retrieval system using FAISS
 def build_faiss_index(texts):
     embeddings = embedder.encode(texts, show_progress_bar=True)
     dimension = embeddings.shape[1]
@@ -113,15 +98,14 @@ def build_faiss_index(texts):
     index.add(embeddings)
     return index
 
-# Context-aware generation
+# Context-aware answer generation
 def generate_answer(query, context):
     prompt = f"""<|user|>
-    Using this context: {context}
-    Answer the question precisely and truthfully. If unsure, say "I don't know".
-    Question: {query}
-    <|assistant|>
-    """
-    
+Using this context: {context}
+Answer the question precisely and truthfully. If unsure, say "I don't know".
+Question: {query}
+<|assistant|>
+"""
     inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True).to(DEVICE)
     outputs = model.generate(
         **inputs,
@@ -131,12 +115,14 @@ def generate_answer(query, context):
         repetition_penalty=1.2,
         do_sample=True
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
+    answer_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer_text.split("<|assistant|>")[-1].strip()
 
-# Streamlit UI
+# Streamlit UI setup
 st.set_page_config(page_title="📚 Smart Book Analyst", layout="wide")
 st.title("📚 AI-Powered Book Analysis System")
 
+# File upload
 uploaded_file = st.file_uploader("Upload book (PDF or TXT)", type=["pdf", "txt"])
 
 if uploaded_file:
@@ -158,7 +144,8 @@ if uploaded_file:
         except Exception as e:
             st.error(f"Processing failed: {str(e)}")
 
-if st.session_state.index:
+# Query interface
+if "index" in st.session_state and st.session_state.index is not None:
     query = st.text_input("Ask about the book:")
     if query:
         with st.spinner("🔍 Searching for answers..."):
@@ -176,4 +163,4 @@ if st.session_state.index:
                 st.caption("Retrieved context confidence: {:.2f}".format(distances[0][0]))
                 
             except Exception as e:
-                st.error(f"Query failed: {str(e)}")
+                st.error(f"Query failed: {str(e)}")