Updated main.py

app.py

@@ -5,99 +5,114 @@ from sentence_transformers import SentenceTransformer
 import faiss
 from transformers import pipeline
 
-# Caching heavy resources
+# Load and cache models
 @st.cache_resource
-def load_models():
-    # Embedding model (lightweight)
+def load_resources():
     embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-    # Generative QA model
-    qa_gen = pipeline(
+    # Generative chat model
+    chat_gen = pipeline(
         'text2text-generation',
         model='google/flan-t5-base',
         tokenizer='google/flan-t5-base',
-        device=-1  # CPU
+        device=-1,
     )
-    return embedder, qa_gen
+    return embedder, chat_gen
 
 # Extract text from uploaded file
-def extract_text_from_file(uploaded_file):
-    name = uploaded_file.name.lower()
+def extract_text(uploaded):
+    name = uploaded.name.lower()
     if name.endswith('.pdf'):
-        reader = PdfReader(uploaded_file)
-        text = ''.join(page.extract_text() or '' for page in reader.pages)
-    elif name.endswith('.docx'):
-        doc = docx.Document(uploaded_file)
-        text = '\n'.join(para.text for para in doc.paragraphs)
-    else:
-        text = uploaded_file.getvalue().decode('utf-8', errors='ignore')
-    return text
+        reader = PdfReader(uploaded)
+        return ''.join(page.extract_text() or '' for page in reader.pages)
+    if name.endswith('.docx'):
+        doc = docx.Document(uploaded)
+        return '\n'.join(para.text for para in doc.paragraphs)
+    return uploaded.getvalue().decode('utf-8', errors='ignore')
 
-# Split text into chunks
-def chunk_text(text, chunk_size=500, overlap=50):
+# Chunking helper
+def chunk_text(text, size=500, overlap=50):
     words = text.split()
-    chunks = []
-    start = 0
+    chunks, start = [], 0
     while start < len(words):
-        end = min(start + chunk_size, len(words))
+        end = min(start + size, len(words))
         chunks.append(' '.join(words[start:end]))
-        start += chunk_size - overlap
+        start += size - overlap
     return chunks
 
-# Build FAISS index from chunks
+# Build FAISS index
 @st.cache_resource
-def build_faiss_index(chunks, _embedder):  # underscore avoids hashing
-    embeddings = _embedder.encode(chunks, convert_to_numpy=True)
-    dim = embeddings.shape[1]
+def build_index(chunks, _embedder):  # underscore avoids hashing
+    embs = _embedder.encode(chunks, convert_to_numpy=True)
+    dim = embs.shape[1]
     index = faiss.IndexFlatL2(dim)
-    index.add(embeddings)
+    index.add(embs)
     return index
 
-# Main Streamlit app
+# Compose prompt for chat+RAG
+def make_prompt(system_prompt, context, history, question):
+    prompt = system_prompt + "\n\n" + "Document Context:\n" + context + "\n\n"
+    # append conversation history
+    for msg in history:
+        role, text = msg['role'], msg['text']
+        prompt += f"{role}: {text}\n"
+    prompt += f"User: {question}\nAssistant:"
+    return prompt
+
+# Main app
 def main():
-    st.set_page_config(page_title='📄 RAGbot', layout='wide')
+    st.set_page_config(page_title='📄 RagBot Chat+RAG', layout='wide')
     st.title('🤖 RagBot')
-    st.sidebar.header('Upload Documents')
+    st.sidebar.header('📂 Upload Document')
 
-    # Initialize chat history
+    # Initialize state
     if 'history' not in st.session_state:
-        st.session_state.history = []
+        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
 
-    uploaded = st.sidebar.file_uploader('Upload PDF/DOCX/TXT', type=['pdf', 'docx', 'txt'])
-    if not uploaded:
-        st.info('Please upload a document in the sidebar to begin.')
-        return
-
-    # On first load of a doc, process and index
-    if 'chunks' not in st.session_state or st.session_state.uploaded_name != uploaded.name:
-        text = extract_text_from_file(uploaded)
+    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.qa_gen = load_models()
-        st.session_state.index = build_faiss_index(st.session_state.chunks, st.session_state.embedder)
+        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 = []  # reset history on new doc
+        st.session_state.history = []
 
-    # Display existing chat history
-    for chat in st.session_state.history:
-        with st.chat_message('user'):
-            st.markdown(f"**You:** {chat['question']}")
-        with st.chat_message('assistant'):
-            st.markdown(f"**RagBot:** {chat['answer']}")
+    # 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
+
+    # 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']}")
 
-    # Chat input
-    question = st.chat_input('Ask a question about the document...')
+    # User input
+    question = st.chat_input('Ask anything—general or about the document...')
     if question:
-        # Retrieve top-k relevant chunks
+        # Retrieve relevant context
         q_emb = st.session_state.embedder.encode([question], convert_to_numpy=True)
-        _, I = st.session_state.index.search(q_emb, k=3)
-        context = '\n\n'.join(st.session_state.chunks[i] for i in I[0])
+        _, idxs = st.session_state.index.search(q_emb, k=3)
+        context = '\n\n'.join(st.session_state.chunks[i] for i in idxs[0])
 
-        # Generate answer
-        prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer in detail:"
-        response = st.session_state.qa_gen(prompt, max_new_tokens=200, do_sample=False)
+        # Build and run 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."
+        )
+        prompt = make_prompt(system_prompt, context, st.session_state.history, question)
+        response = st.session_state.chat_gen(prompt, max_new_tokens=200, do_sample=False)
         answer = response[0]['generated_text'].strip()
 
-        # Save & display new messages
-        st.session_state.history.append({'question': question, 'answer': answer})
+        # Record and display
+        st.session_state.history.append({'role':'User','text':question})
+        st.session_state.history.append({'role':'Assistant','text':answer})
         with st.chat_message('user'):
             st.markdown(f"**You:** {question}")
         with st.chat_message('assistant'):