Updated main.py

app.py

@@ -9,7 +9,6 @@ from transformers import pipeline
 @st.cache_resource
 def load_resources():
     embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-    # Generative chat model
     chat_gen = pipeline(
         'text2text-generation',
         model='google/flan-t5-base',
@@ -48,13 +47,14 @@ def build_index(chunks, _embedder):  # underscore avoids hashing
     index.add(embs)
     return index
 
-# Compose prompt for chat+RAG
+# Compose prompt
 def make_prompt(system_prompt, context, history, question):
-    prompt = system_prompt + "\n\n" + "Document Context:\n" + context + "\n\n"
-    # append conversation history
+    prompt = system_prompt + "\n\n"
+    if context:
+        prompt += f"Document Context:\n{context}\n\n"
     for msg in history:
-        role, text = msg['role'], msg['text']
-        prompt += f"{role}: {text}\n"
+        role = 'User' if msg['role']=='User' else 'Assistant'
+        prompt += f"{role}: {msg['text']}\n"
     prompt += f"User: {question}\nAssistant:"
     return prompt
 
@@ -62,55 +62,57 @@ def make_prompt(system_prompt, context, history, question):
 def main():
     st.set_page_config(page_title='📄 RagBot Chat+RAG', layout='wide')
     st.title('🤖 RagBot')
-    st.sidebar.header('📂 Upload Document')
+    st.sidebar.header('📂 Optional: Upload Document')
 
     # Initialize state
     if 'history' not in st.session_state:
-        st.session_state.history = []  # list of {'role': 'User|Assistant', 'text': ...}
+        st.session_state.history = []  # list of {'role': 'User'|'Assistant', 'text': ...}
     if 'chunks' not in st.session_state:
         st.session_state.chunks = []
     if 'index' not in st.session_state:
         st.session_state.index = None
 
+    # Document upload
     uploaded = st.sidebar.file_uploader('Upload PDF, DOCX or TXT', type=['pdf','docx','txt'])
     if uploaded and (st.session_state.get('uploaded_name') != uploaded.name):
-        # New document: extract, chunk, index, reset
         text = extract_text(uploaded)
         st.session_state.chunks = chunk_text(text)
         st.session_state.embedder, st.session_state.chat_gen = load_resources()
         st.session_state.index = build_index(st.session_state.chunks, st.session_state.embedder)
         st.session_state.uploaded_name = uploaded.name
-        st.session_state.history = []
-
-    # If no doc yet, ask to upload
-    if st.session_state.index is None:
-        st.info('Please upload a document in the sidebar to start.')
-        return
+        st.session_state.history = []  # reset conversation
+    # Load models if not loaded
+    if 'embedder' not in st.session_state or 'chat_gen' not in st.session_state:
+        st.session_state.embedder, st.session_state.chat_gen = load_resources()
 
     # Display chat history
     for msg in st.session_state.history:
         with st.chat_message('user' if msg['role']=='User' else 'assistant'):
             st.markdown(f"**{msg['role']}:** {msg['text']}")
 
-    # User input
-    question = st.chat_input('Ask anything—general or about the document...')
+    # Chat input always available
+    question = st.chat_input('Ask a question—general or document-specific...')
     if question:
-        # Retrieve relevant context
-        q_emb = st.session_state.embedder.encode([question], convert_to_numpy=True)
-        _, idxs = st.session_state.index.search(q_emb, k=3)
-        context = '\n\n'.join(st.session_state.chunks[i] for i in idxs[0])
+        # Retrieve context if index exists
+        context = ''
+        if st.session_state.index is not None:
+            q_emb = st.session_state.embedder.encode([question], convert_to_numpy=True)
+            _, idxs = st.session_state.index.search(q_emb, k=3)
+            context = '\n\n'.join(st.session_state.chunks[i] for i in idxs[0])
 
-        # Build and run prompt
+        # Build prompt
         system_prompt = (
             "You are RagBot, an AI assistant. "
-            "Use the provided document context to answer specific questions, "
-            "but also leverage your general knowledge for broader queries."
+            "Use the provided document context for specific questions, "
+            "and your general knowledge for everything else."
         )
         prompt = make_prompt(system_prompt, context, st.session_state.history, question)
+
+        # Generate answer
         response = st.session_state.chat_gen(prompt, max_new_tokens=200, do_sample=False)
         answer = response[0]['generated_text'].strip()
 
-        # Record and display
+        # Save & display
         st.session_state.history.append({'role':'User','text':question})
         st.session_state.history.append({'role':'Assistant','text':answer})
         with st.chat_message('user'):