Update app.py

app.py

@@ -5,7 +5,6 @@ import faiss
 import numpy as np
 from sentence_transformers import SentenceTransformer
 import PyPDF2
-import os
 
 # Model Setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -22,6 +21,8 @@ embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 dimension = 384  # Embedding size for MiniLM
 index = faiss.IndexFlatL2(dimension)
 docs = []  # Store document texts
+summary = ""  # Store book summary
+
 
 # Function to extract text from PDF
 def extract_text_from_pdf(uploaded_file):
@@ -29,9 +30,10 @@ def extract_text_from_pdf(uploaded_file):
     text = "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
     return text
 
-# Function to process uploaded documents
+
+# Function to process uploaded documents and generate summary
 def process_documents(files):
-    global docs, index
+    global docs, index, summary
     docs = []
     
     for file in files:
@@ -44,15 +46,38 @@ def process_documents(files):
     
     embeddings = embedding_model.encode(docs)
     index.add(np.array(embeddings))
+    
+    # Generate summary after processing documents
+    summary = generate_summary("\n".join(docs))
+
+
+# Function to generate a book summary
+def generate_summary(text):
+    chat = [
+        {"role": "system", "content": "You are a helpful AI that summarizes books."},
+        {"role": "user", "content": f"Summarize this book in a short paragraph:\n{text[:4000]}"}  # Limiting input size
+    ]
+    chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+
+    input_tokens = tokenizer(chat, return_tensors="pt").to(device)
+    output = model.generate(**input_tokens, max_new_tokens=300)
+    
+    return tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+
 
 # Function to retrieve relevant context
 def retrieve_context(query):
+    if index.ntotal == 0:
+        return "No documents available. Please upload files first."
+    
     query_embedding = embedding_model.encode([query])
     distances, indices = index.search(np.array(query_embedding), k=1)
     
-    if len(indices) > 0 and indices[0][0] < len(docs):
-        return docs[indices[0][0]]
-    return "No relevant context found."
+    if len(indices) == 0 or indices[0][0] >= len(docs):
+        return "No relevant context found."
+
+    return docs[indices[0][0]]
+
 
 # Function to generate response using IBM Granite
 def generate_response(query, context):
@@ -64,27 +89,32 @@ def generate_response(query, context):
 
     input_tokens = tokenizer(chat, return_tensors="pt").to(device)
     output = model.generate(**input_tokens, max_new_tokens=200)
+    
     return tokenizer.batch_decode(output, skip_special_tokens=True)[0]
 
+
 # Streamlit UI
-st.set_page_config(page_title="📖 Smart Study", page_icon="🤖")
-st.title("📖 Q&A using IBM Granite")
-st.subheader("Upload documents and ask questions!")
+st.set_page_config(page_title="📖 AI Book Assistant", page_icon="📚")
+st.title("📖 AI-Powered Book Assistant")
+st.subheader("Upload a book and get its summary or ask questions!")
 
-uploaded_files = st.file_uploader("Upload PDFs or TXT files", accept_multiple_files=True)
+uploaded_files = st.file_uploader("Upload a book (PDF or TXT)", accept_multiple_files=False)
 
 if uploaded_files:
-    with st.spinner("Processing documents..."):
-        process_documents(uploaded_files)
-    st.success("Documents uploaded and indexed!")
+    with st.spinner("Processing book and generating summary..."):
+        process_documents([uploaded_files])
+    st.success("Book uploaded and processed!")
 
-if query:
-    if index.ntotal == 0:  # Ensure documents are indexed before querying
-        st.warning("Please upload and process documents first!")
+    st.markdown("### 📚 Book Summary:")
+    st.write(summary)
+
+query = st.text_input("Ask a question about the book:")
+if st.button("Get Answer"):
+    if index.ntotal == 0:
+        st.warning("Please upload a book first!")
     else:
         with st.spinner("Retrieving and generating response..."):
             context = retrieve_context(query)
             response = generate_response(query, context)
             st.markdown("### 🤖 Answer:")
             st.write(response)
-