Update app.py

app.py

@@ -4,7 +4,7 @@ st.set_page_config(page_title="RAG Book Analyzer", layout="wide")  # Must be the
 import torch
 import numpy as np
 import faiss
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from sentence_transformers import SentenceTransformer
 import fitz  # PyMuPDF for PDF extraction
 import docx2txt  # For DOCX extraction
@@ -13,7 +13,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 # ------------------------
 # Configuration
 # ------------------------
-MODEL_NAME = "ibm-granite/granite-3.1-1b-a400m-instruct"
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
 EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
 CHUNK_SIZE = 512
 CHUNK_OVERLAP = 64
@@ -25,19 +25,13 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 @st.cache_resource
 def load_models():
     try:
-        tokenizer = AutoTokenizer.from_pretrained(
-            MODEL_NAME,
-            trust_remote_code=True,
-            revision="main"
-        )
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            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()
+        )
         embedder = SentenceTransformer(EMBED_MODEL, device=DEVICE)
         return tokenizer, model, embedder
     except Exception as e:
@@ -79,10 +73,9 @@ def extract_text(file):
         return ""
 
 def build_index(chunks):
-    embeddings = embedder.encode(chunks, show_progress_bar=True)
+    embeddings = embedder.encode(chunks, show_progress_bar=False)
     dimension = embeddings.shape[1]
-    index = faiss.IndexFlatIP(dimension)
-    faiss.normalize_L2(embeddings)
+    index = faiss.IndexFlatL2(dimension)
     index.add(embeddings)
     return index
 
@@ -90,36 +83,38 @@ def build_index(chunks):
 # 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"
+    # Create prompt with Mistral format
+    prompt = f"<s>[INST] Summarize this book in a concise paragraph: {text[:3000]} [/INST]"
     inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-    outputs = model.generate(**inputs, max_new_tokens=300, temperature=0.5)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=300,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True
+    )
     summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Remove any markers and extra lines; return the first non-empty paragraph.
-    summary = summary.replace("<|assistant|>", "").strip()
-    paragraphs = [p.strip() for p in summary.split("\n") if p.strip()]
-    return paragraphs[0] if paragraphs else summary
+    return summary.split("[/INST]")[-1].strip()
 
 def generate_answer(query, context):
-    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)
+    # Create prompt with Mistral format
+    prompt = f"<s>[INST] Answer this question based on the context. If unsure, say 'I don't know'.\n\nQuestion: {query}\nContext: {context} [/INST]"
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
     outputs = model.generate(
         **inputs,
         max_new_tokens=300,
-        temperature=0.4,
+        temperature=0.5,
         top_p=0.9,
         repetition_penalty=1.2,
         do_sample=True
     )
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = answer.replace("<|assistant|>", "").strip()
-    paragraphs = [p.strip() for p in answer.split("\n") if p.strip()]
-    return paragraphs[0] if paragraphs else answer
+    return answer.split("[/INST]")[-1].strip()
 
 # ------------------------
 # Streamlit UI
 # ------------------------
-st.title("RAG-Based Book Analyzer")
+st.title("📚 RAG-Based Book Analyzer")
 st.write("Upload a book (PDF, TXT, DOCX) to get a summary and ask questions about its content.")
 
 uploaded_file = st.file_uploader("Upload File", type=["pdf", "txt", "docx"])
@@ -127,11 +122,12 @@ uploaded_file = st.file_uploader("Upload File", type=["pdf", "txt", "docx"])
 if uploaded_file:
     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)
+        st.success("✅ File successfully processed!")
+        
+        with st.spinner("Generating summary..."):
+            summary = generate_summary(text)
+            st.markdown("### Book Summary")
+            st.info(summary)
         
         # Process text into chunks and build FAISS index
         chunks = split_text(text)
@@ -139,15 +135,19 @@ if uploaded_file:
         st.session_state.chunks = chunks
         st.session_state.index = index
         
-        st.markdown("### Ask a Question about the Book:")
-        query = st.text_input("Your Question:")
+        st.markdown("### ❓ Ask a Question about the Book")
+        query = st.text_input("Enter your question:")
         if query:
-            # Retrieve top 3 relevant chunks as context
-            query_embedding = embedder.encode([query])
-            faiss.normalize_L2(query_embedding)
-            distances, indices = st.session_state.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)
+            with st.spinner("Searching for answers..."):
+                # Retrieve top 3 relevant chunks as context
+                query_embedding = embedder.encode([query])
+                distances, indices = st.session_state.index.search(query_embedding, k=3)
+                retrieved_chunks = [st.session_state.chunks[i] for i in indices[0] if i < len(st.session_state.chunks)]
+                context = "\n\n".join(retrieved_chunks)
+                
+                answer = generate_answer(query, context)
+                st.markdown("### 💬 Answer")
+                st.success(answer)
+                
+                with st.expander("See context used"):
+                    st.write(context)