Updated main.py

app.py

@@ -14,6 +14,9 @@ def load_resources():
         model='google/flan-t5-base',
         tokenizer='google/flan-t5-base',
         device=-1,
+        # enforce deterministic decoding and low temperature to reduce hallucinations
+        do_sample=False,
+        temperature=0.0,
     )
     return embedder, chat_gen
 
@@ -66,7 +69,7 @@ def main():
 
     # Initialize state
     if 'history' not in st.session_state:
-        st.session_state.history = []  # list of {'role': 'User'|'Assistant', 'text': ...}
+        st.session_state.history = []
     if 'chunks' not in st.session_state:
         st.session_state.chunks = []
     if 'index' not in st.session_state:
@@ -80,8 +83,8 @@ def main():
         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 conversation
-    # Load models if not loaded
+        st.session_state.history = []
+    # Load models if missing
     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()
 
@@ -93,23 +96,25 @@ def main():
     # Chat input always available
     question = st.chat_input('Ask a question—general or document-specific...')
     if question:
-        # Retrieve context if index exists
+        # Retrieve context
         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 prompt
+        # Build prompt with hallucination guard
         system_prompt = (
             "You are RagBot, an AI assistant. "
-            "Use the provided document context for specific questions, "
-            "and your general knowledge for everything else."
+            "You must ONLY use the document context provided to answer document-specific questions. "
+            "If the answer is not contained in the context, respond with: "
+            "\"I’m sorry, I don’t know based on the document.\" "
+            "For general knowledge questions, answer using your training knowledge without hallucinating."
         )
         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)
+        response = st.session_state.chat_gen(prompt, max_new_tokens=200)
         answer = response[0]['generated_text'].strip()
 
         # Save & display