Update app.py

app.py

@@ -1,87 +1,125 @@
-import os
 import gradio as gr
-
-# ✅ New LangChain Hugging Face imports
-from langchain_huggingface import HuggingFaceEndpoint
+from PyPDF2 import PdfReader
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.chains import ConversationalRetrievalChain
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.docstore.document import Document
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# -----------------------
+# 1️⃣ Load PDF & Split
+# -----------------------
+pdf_path = "chimera.pdf"
 
-# --- 1️⃣ Load your PDF ---
-loader = PyPDFLoader("chimera.pdf")
-documents = loader.load()
+reader = PdfReader(pdf_path)
+evidences = [page.extract_text() for page in reader.pages if page.extract_text()]
 
-# --- 2️⃣ Split into chunks ---
+# Split each evidence into chunks
+docs = [Document(page_content=text) for text in evidences]
 text_splitter = CharacterTextSplitter(chunk_size=800, chunk_overlap=100)
-texts = text_splitter.split_documents(documents)
+texts = text_splitter.split_documents(docs)
 
-# --- 3️⃣ Create embeddings + FAISS vector store ---
+# Embeddings & FAISS
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 db = FAISS.from_documents(texts, embeddings)
 retriever = db.as_retriever(search_kwargs={"k": 3})
 
-# --- 4️⃣ Hugging Face Endpoint setup ---
-hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-if hf_token is None:
-    raise ValueError(
-        "HUGGINGFACEHUB_API_TOKEN not found! Add it in Space Settings → Secrets."
-    )
+# -----------------------
+# 2️⃣ Local LLM
+# -----------------------
+llm_model_name = "google/flan-t5-small"
+tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)
 
-llm = HuggingFaceEndpoint(
-    endpoint_url="https://api-inference.huggingface.co/models/google/flan-t5-base",
-    headers={"Authorization": f"Bearer {hf_token}"},
-    model_kwargs={"temperature": 0}
-)
+def generate_answer(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=150)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# --- 5️⃣ Build conversational chain ---
-qa = ConversationalRetrievalChain.from_llm(llm, retriever=retriever)
+# -----------------------
+# 3️⃣ Evidence Navigation & Chat
+# -----------------------
+def enter_case():
+    msg = f"Evidence 1 of {len(evidences)}:\n\n{evidences[0]}"
+    return msg, 0, 0, gr.update(interactive=True), gr.update(interactive=False), gr.update(visible=False)
 
-chat_history = []
+def next_evidence(idx):
+    if idx + 1 < len(evidences):
+        idx += 1
+        return f"Evidence {idx+1} of {len(evidences)}:\n\n{evidences[idx]}", idx, 0, gr.update(interactive=True), gr.update(interactive=False), gr.update(visible=False)
+    return "All evidences reviewed. Investigation completed.", idx, 0, gr.update(interactive=False), gr.update(interactive=False), gr.update(visible=False)
 
-# --- 6️⃣ Respond function ---
-def respond(message, history):
-    history = history[-6:]  # keep last 3 exchanges
-    result = qa({"question": message, "chat_history": history})
-    history.append((message, result["answer"]))
-    return history, history
+def ask_question(message, history, idx, qcount):
+    if qcount >= 3:
+        return history, qcount, gr.update(interactive=False), gr.update(interactive=True)
+    
+    relevant_docs = retriever.get_relevant_documents(message)
+    context = "\n".join([doc.page_content for doc in relevant_docs])
+    prompt = f"Context: {context}\n\nQuestion: {message}\nAnswer:"
+    answer = generate_answer(prompt)
+    
+    history = history or []
+    history.append((message, answer))
+    qcount += 1
+    
+    disable_input = gr.update(interactive=(qcount < 3))
+    enable_next = gr.update(interactive=(qcount >= 3))
+    
+    return history, qcount, disable_input, enable_next
 
-# --- 7️⃣ Gradio UI ---
+# -----------------------
+# 4️⃣ Gradio UI
+# -----------------------
 with gr.Blocks() as demo:
+    # Warning Message
     with gr.Column():
-        warning_text = gr.HTML(
-            "<div style='background-color:black;color:white;padding:20px;'>⚠ WARNING: Investigative Simulation ⚠<br>Are you ready?</div>"
+        warning_msg = gr.Markdown(
+            """
+            ⚠ **WARNING — INVESTIGATIVE SIMULATION** ⚠  
+            You are about to enter The Chimera Case, a high-stakes investigation into Innovate Future Labs (IFL)
+            and Project Chimera.  
+            There are 11 pieces of evidence. For each evidence, you can ask **only 3 questions**.  
+            Total questions allowed: 33.  
+            Treat every claim as unverified until verified by evidence.  
+            Are you ready to proceed?  
+            """,
         )
         enter_btn = gr.Button("Enter the Case")
-        exit_btn = gr.Button("Exit")
-        chatbot = gr.Chatbot(visible=False)
-        user_input = gr.Textbox(placeholder="Type here...", visible=False)
-        submit_btn = gr.Button("Send", visible=False)
-
-    def enter_case():
-        return (
-            gr.update(visible=True),  # chatbot
-            gr.update(visible=True),  # user_input
-            gr.update(visible=True),  # submit_btn
-            gr.update(value=""),      # hide warning
-            gr.update(visible=False), # hide enter_btn
-            gr.update(visible=False)  # hide exit_btn
-        )
-
-    def exit_case():
-        return (
-            gr.update(value="Session ended."),
-            gr.update(visible=False),
-            gr.update(visible=False),
-            gr.update(visible=False),
-            gr.update(visible=False),
-            gr.update(visible=False)
-        )
-
-    enter_btn.click(enter_case, inputs=None, outputs=[chatbot, user_input, submit_btn, warning_text, enter_btn, exit_btn])
-    exit_btn.click(exit_case, inputs=None, outputs=[warning_text, chatbot, user_input, submit_btn, enter_btn, exit_btn])
-    submit_btn.click(respond, inputs=[user_input, chatbot], outputs=[chatbot, chatbot])
+    
+    # Evidence display
+    evidence_box = gr.Textbox(label="Evidence", interactive=False, lines=10, visible=False)
+    next_btn = gr.Button("Next Evidence", interactive=False, visible=False)
+    
+    # Chatbot
+    chatbot = gr.Chatbot()
+    user_input = gr.Textbox(placeholder="Ask a question about this evidence...", interactive=False)
+    submit_btn = gr.Button("Send", interactive=False)
+    
+    state_idx = gr.State(value=0)
+    q_count = gr.State(value=0)
+    
+    # -----------------------
+    # Button Actions
+    # -----------------------
+    enter_btn.click(
+        enter_case, 
+        outputs=[evidence_box, state_idx, q_count, user_input, next_btn, enter_btn]
+    )
+    
+    next_btn.click(
+        next_evidence, 
+        inputs=[state_idx], 
+        outputs=[evidence_box, state_idx, q_count, user_input, next_btn, enter_btn]
+    )
+    
+    submit_btn.click(
+        ask_question, 
+        inputs=[user_input, chatbot, state_idx, q_count], 
+        outputs=[chatbot, q_count, user_input, next_btn]
+    )
 
+# -----------------------
+# 5️⃣ Launch
+# -----------------------
 if __name__ == "__main__":
-    demo.launch(share=True, enable_queue=True)
+    demo.launch()