Update app.py

app.py

@@ -33,56 +33,32 @@ retriever = get_retriever()
 pipe = pipeline(
     "text-generation",
     model="google/flan-t5-base",   # ✅ smaller + CPU friendly
-    max_new_tokens=512,
-    temperature=0.3
+    max_new_tokens=256,  # Reduced from 512 to fit within context
+    temperature=0.3,
+    do_sample=True,
+    pad_token_id=0  # Add padding token
 )
 llm = HuggingFacePipeline(pipeline=pipe)
 
 
 # =====================================================
-# Prompts
+# Prompts (shortened to reduce token usage)
 # =====================================================
-english_system_prompt = """
-You are a Nigerian Legal AI Assistant specialized in Nigerian law. You have deep knowledge of:
-- Nigerian Constitution 1999
-- Labour Act and Employment Laws
-- Nigeria Data Protection Act
-- Federal Competition and Consumer Protection Act (FCCPA)
-
-PERSONALITY: Professional but approachable, uses Nigerian legal terminology, understands local context.
-
-RESPONSE STYLE:
-- Start with direct answer to the question
-- Quote specific sections/articles when available
-- Explain in simple terms what the law means
-- Always include disclaimer: "⚠️ This is not legal advice. Please consult a qualified lawyer for specific issues."
-- Use Nigerian English expressions naturally (but not forced)
-
-CONVERSATION MEMORY: Remember previous questions in this chat to provide contextual follow-ups.
+english_system_prompt = """You are a Nigerian Legal AI Assistant. Provide direct answers about Nigerian law with relevant sections/articles. Always end with: "⚠️ This is not legal advice. Consult a qualified lawyer."
 """
 
-pidgin_system_prompt = """
-You be Nigerian Legal AI Assistant wey sabi Nigerian law well well. You get knowledge of:
-- Nigerian Constitution 1999
-- Labour Act and Employment Laws
-- Nigeria Data Protection Act
-- Federal Competition and Consumer Protection Act (FCCPA)
-
-PERSONALITY: Friendly, approachable, dey use Naija way of talk but still correct for legal matter.
-
-RESPONSE STYLE:
-- Start with direct answer
-- Mention the exact section/article if available
-- Explain am for clear Pidgin wey anybody fit understand
-- Always add disclaimer: "⚠️ No be legal advice o, abeg meet lawyer if matter serious."
-- Remember wetin dem don ask before, make conversation flow well.
+pidgin_system_prompt = """You be Nigerian Legal AI Assistant. Give direct answer about Nigerian law with correct section/article. Always end with: "⚠️ No be legal advice o, abeg meet lawyer."
 """
 
 
 # =====================================================
-# Conversational QA Chain
+# Conversational QA Chain with fixed memory
 # =====================================================
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+memory = ConversationBufferMemory(
+    memory_key="chat_history", 
+    return_messages=True,
+    output_key="answer"  # Fix: specify which output to store in memory
+)
 
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm=llm,
@@ -93,31 +69,55 @@ qa_chain = ConversationalRetrievalChain.from_llm(
 
 
 # =====================================================
-# Chat function
+# Chat function with better token management
 # =====================================================
 def answer_question(user_input, lang_choice, history=[]):
-    # Pick system prompt
-    if lang_choice == "pidgin":
-        system_prompt = pidgin_system_prompt
-    else:
-        system_prompt = english_system_prompt
-
-    # Run QA
-    result = qa_chain.invoke({"question": f"{system_prompt}\n\nUser: {user_input}"})
-    answer = result["answer"]
-
-    # Collect sources (with sections)
-    sources = []
-    for doc in result["source_documents"]:
-        section = doc.metadata.get("section", "Unknown Section")
-        source = doc.metadata.get("source", "Unknown Document").replace(".pdf", "")
-        sources.append(f"[{section}] from {source}")
-
-    if sources:
-        answer += "\n\n📚 Sources:\n" + "\n".join(sources)
-
-    history.append(("You: " + user_input, "Bot: " + answer))
-    return history, history
+    try:
+        # Pick system prompt
+        if lang_choice == "pidgin":
+            system_prompt = pidgin_system_prompt
+        else:
+            system_prompt = english_system_prompt
+
+        # Truncate user input if too long
+        max_input_length = 200  # Limit user input length
+        if len(user_input) > max_input_length:
+            user_input = user_input[:max_input_length] + "..."
+
+        # Create shorter question format
+        question = f"{system_prompt}\nQ: {user_input}"
+        
+        # Run QA
+        result = qa_chain.invoke({"question": question})
+        answer = result["answer"]
+
+        # Collect sources (with sections) - limit to top 3
+        sources = []
+        for doc in result["source_documents"][:3]:  # Limit to top 3 sources
+            section = doc.metadata.get("section", "Unknown Section")
+            source = doc.metadata.get("source", "Unknown Document").replace(".pdf", "")
+            sources.append(f"[{section}] from {source}")
+
+        if sources:
+            answer += "\n\n📚 Sources:\n" + "\n".join(sources)
+
+        # Truncate answer if too long
+        max_answer_length = 800
+        if len(answer) > max_answer_length:
+            answer = answer[:max_answer_length] + "...\n\n⚠️ Response truncated due to length limits."
+
+        history.append(("You: " + user_input, "Bot: " + answer))
+        
+        # Limit history to last 5 exchanges to prevent memory overflow
+        if len(history) > 5:
+            history = history[-5:]
+            
+        return history, history
+
+    except Exception as e:
+        error_msg = f"Sorry, I encountered an error: {str(e)[:100]}..."
+        history.append(("You: " + user_input, "Bot: " + error_msg))
+        return history, history
 
 
 # =====================================================
@@ -129,16 +129,24 @@ with gr.Blocks(css=".gradio-container {max-width: 800px !important}") as demo:
     with gr.Row():
         with gr.Column(scale=4):
             chatbot = gr.Chatbot(label="Chat with Legal AI", height=500)
-            msg = gr.Textbox(label="Ask your question here...")
+            msg = gr.Textbox(
+                label="Ask your question here...", 
+                placeholder="e.g., What are the rights of employees in Nigeria?",
+                max_lines=3
+            )
             lang_choice = gr.Radio(["english", "pidgin"], value="english", label="Language")
             clear = gr.Button("Clear Chat")
 
     state = gr.State([])
 
     def reset():
+        # Clear memory as well
+        global memory
+        memory.clear()
         return [], []
 
     msg.submit(answer_question, [msg, lang_choice, state], [chatbot, state])
+    msg.submit(lambda: "", None, msg)  # Clear input after submit
     clear.click(reset, None, [chatbot, state])
 
-demo.launch()
+demo.launch()