Update app.py

app.py

@@ -5,19 +5,24 @@ import faiss
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 import fitz  # PyMuPDF for PDF extraction
+import docx2txt  # For DOCX 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"
 
+# ------------------------
+# Model Loading with Caching
+# ------------------------
 @st.cache_resource
 def load_models():
     try:
-        # Load tokenizer and generative model with trust_remote_code enabled
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
             trust_remote_code=True,
@@ -25,25 +30,24 @@ def load_models():
         )
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            device_map="auto" if DEVICE == "cuda" else None,
-            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
             trust_remote_code=True,
             revision="main",
+            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
-    
     except Exception as e:
         st.error(f"Model loading failed: {str(e)}")
         st.stop()
 
 tokenizer, model, embedder = load_models()
 
-# Improved text processing: splits text into chunks
-def process_text(text):
+# ------------------------
+# Text Processing Functions
+# ------------------------
+def split_text(text):
     splitter = RecursiveCharacterTextSplitter(
         chunk_size=CHUNK_SIZE,
         chunk_overlap=CHUNK_OVERLAP,
@@ -51,61 +55,48 @@ def process_text(text):
     )
     return splitter.split_text(text)
 
-# Enhanced PDF extraction using PyMuPDF
-def extract_pdf_text(uploaded_file):
-    try:
-        doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
-        return "\n".join([page.get_text() for page in doc])
-    except Exception as e:
-        st.error(f"PDF extraction error: {str(e)}")
+def extract_text(file):
+    file_type = file.type
+    if file_type == "application/pdf":
+        try:
+            doc = fitz.open(stream=file.read(), filetype="pdf")
+            return "\n".join([page.get_text() for page in doc])
+        except Exception as e:
+            st.error("Error processing PDF: " + str(e))
+            return ""
+    elif file_type == "text/plain":
+        return file.read().decode("utf-8")
+    elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+        try:
+            return docx2txt.process(file)
+        except Exception as e:
+            st.error("Error processing DOCX: " + str(e))
+            return ""
+    else:
+        st.error("Unsupported file type: " + file_type)
         return ""
 
-# Multi-step summarization
-def generate_summary(text):
-    chunks = process_text(text)[:10]  # Use first 10 chunks for summary
-    summaries = []
-    
-    for chunk in chunks:
-        prompt = f"""<|user|>
-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)
-        summary_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        summaries.append(summary_text)
-    
-    # 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:"""
-    
-    inputs = tokenizer(final_prompt, return_tensors="pt").to(DEVICE)
-    outputs = model.generate(**inputs, max_new_tokens=500, temperature=0.5)
-    full_summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return full_summary.split(":")[-1].strip()
-
-# Enhanced retrieval system using FAISS
-def build_faiss_index(texts):
-    embeddings = embedder.encode(texts, show_progress_bar=True)
+def build_index(chunks):
+    embeddings = embedder.encode(chunks, show_progress_bar=True)
     dimension = embeddings.shape[1]
     index = faiss.IndexFlatIP(dimension)
     faiss.normalize_L2(embeddings)
     index.add(embeddings)
     return index
 
-# Context-aware answer generation
+# ------------------------
+# Summarization and Q&A Functions
+# ------------------------
+def generate_summary(text):
+    # Limit input text to avoid long sequences
+    prompt = f"<|user|>\nSummarize the following book in a concise and informative paragraph:\n\n{text[:4000]}\n<|assistant|>\n"
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    outputs = model.generate(**inputs, max_new_tokens=300, temperature=0.5)
+    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return summary.split("<|assistant|>")[-1].strip() if "<|assistant|>" in summary else summary.strip()
+
 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|>
-"""
+    prompt = f"<|user|>\nUsing the context below, answer the following question precisely. If unsure, say 'I don't know'.\n\nContext: {context}\n\nQuestion: {query}\n<|assistant|>\n"
     inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True).to(DEVICE)
     outputs = model.generate(
         **inputs,
@@ -115,52 +106,42 @@ Question: {query}
         repetition_penalty=1.2,
         do_sample=True
     )
-    answer_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer_text.split("<|assistant|>")[-1].strip()
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer.split("<|assistant|>")[-1].strip() if "<|assistant|>" in answer else answer.strip()
 
-# Streamlit UI setup
-st.set_page_config(page_title="📚 Smart Book Analyst", layout="wide")
-st.title("📚 AI-Powered Book Analysis System")
+# ------------------------
+# Streamlit UI
+# ------------------------
+st.set_page_config(page_title="RAG Book Analyzer", layout="wide")
+st.title("RAG-Based Book Analyzer")
+st.write("Upload a book (PDF, TXT, DOCX) to get a summary and ask questions about its content.")
 
-# File upload
-uploaded_file = st.file_uploader("Upload book (PDF or TXT)", type=["pdf", "txt"])
+uploaded_file = st.file_uploader("Upload File", type=["pdf", "txt", "docx"])
 
 if uploaded_file:
-    with st.spinner("📖 Analyzing book content..."):
-        try:
-            if uploaded_file.type == "application/pdf":
-                text = extract_pdf_text(uploaded_file)
-            else:
-                text = uploaded_file.read().decode()
-            
-            chunks = process_text(text)
-            st.session_state.docs = chunks
-            st.session_state.index = build_faiss_index(chunks)
-            
-            with st.expander("📝 Book Summary", expanded=True):
-                summary = generate_summary(text)
-                st.write(summary)
-                
-        except Exception as e:
-            st.error(f"Processing failed: {str(e)}")
-
-# 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..."):
-            try:
-                # Retrieve top 3 relevant chunks
-                query_embed = embedder.encode([query])
-                faiss.normalize_L2(query_embed)
-                distances, indices = st.session_state.index.search(query_embed, k=3)
-                
-                context = "\n".join([st.session_state.docs[i] for i in indices[0]])
-                answer = generate_answer(query, context)
-                
-                st.subheader("Answer")
-                st.markdown(f"```\n{answer}\n```")
-                st.caption("Retrieved context confidence: {:.2f}".format(distances[0][0]))
-                
-            except Exception as e:
-                st.error(f"Query failed: {str(e)}")
+    text = extract_text(uploaded_file)
+    if text:
+        st.success("File successfully processed!")
+        st.write("Generating summary...")
+        summary = generate_summary(text)
+        st.markdown("### Book Summary")
+        st.write(summary)
+        
+        # Process text into chunks and build FAISS index
+        chunks = split_text(text)
+        index = build_index(chunks)
+        st.session_state.chunks = chunks
+        st.session_state.index = index
+        
+        st.markdown("### Ask a Question about the Book:")
+        query = st.text_input("Your Question:")
+        if query:
+            # Retrieve top 3 relevant chunks as context
+            query_embedding = embedder.encode([query])
+            faiss.normalize_L2(query_embedding)
+            distances, indices = index.search(query_embedding, k=3)
+            retrieved_chunks = [chunks[i] for i in indices[0] if i < len(chunks)]
+            context = "\n".join(retrieved_chunks)
+            answer = generate_answer(query, context)
+            st.markdown("### Answer")
+            st.write(answer)