Add application file

app.py

@@ -5,11 +5,13 @@ import pickle
 from sentence_transformers import SentenceTransformer
 from transformers import pipeline
 
-# File names for saved data
+# File names for saved PDF-based data
 INDEX_FILE = "faiss_index.index"
 CHUNKS_FILE = "chunks.pkl"
+
+# Models
 EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
-QA_MODEL_NAME = "deepset/roberta-large-squad2"  # You can change this to any Hugging Face QA model
+QA_MODEL_NAME = "deepset/roberta-large-squad2"  # More powerful QA model
 
 @st.cache_resource
 def load_index_and_chunks():
@@ -24,13 +26,67 @@ def load_embedding_model():
 
 @st.cache_resource
 def load_qa_pipeline():
-    # This QA pipeline expects a question and a context
     return pipeline("question-answering", model=QA_MODEL_NAME, tokenizer=QA_MODEL_NAME)
 
+# Curated Q/A pairs for enhanced context
+curated_qa_pairs = [
+    {
+        "question": "What is Adagrasib (Krazati®) used for?",
+        "answer": "Adagrasib (Krazati®) is used for the treatment of adult patients with advanced non-small cell lung cancer (NSCLC) with KRAS G12C mutation and disease progression after at least one prior systemic therapy."
+    },
+    {
+        "question": "What are the comparator treatments for Adagrasib in the PICO exercises?",
+        "answer": "In PICOs 1 and 2, the comparator treatments are Sotorasib monotherapy and Platinum-based doublet chemotherapy. In PICO 3, the comparator treatment is Docetaxel monotherapy. In PICO 4, the comparator treatment is Immune checkpoint inhibitor monotherapy."
+    },
+    {
+        "question": "What are the effectiveness measures used in the PICO framework?",
+        "answer": "Effectiveness measures in the PICO framework typically include outcomes such as overall response rate, progression-free survival, and overall survival. In these PICOs, Adagrasib is used as monotherapy in PICOs 1 and 2, and as monotherapy or in combination with platinum-based chemotherapy in PICO 3, while in PICO 4 it may be used as monotherapy or in combination therapy."
+    },
+    {
+        "question": "What adverse events were monitored in the PICO studies?",
+        "answer": "The adverse events monitored in the PICO studies include a range of treatment-related side effects, though the provided context does not detail them fully. More detailed study data would be required for a complete list."
+    },
+    {
+        "question": "How does Adagrasib compare to immune checkpoint inhibitors?",
+        "answer": "The analysis suggests that comparisons between Adagrasib and immune checkpoint inhibitors (or Docetaxel) should account for patients' prior treatment lines. Direct comparisons may not be relevant in certain subgroups, especially in patients with only one prior line of systemic therapy."
+    }
+]
+
+def get_curated_context(query, curated_qa, embed_model):
+    """
+    Retrieve the most relevant curated Q/A pair based on the user's query.
+    Returns a formatted string if the similarity (based on L2 distance) is below a threshold.
+    """
+    curated_questions = [qa["question"] for qa in curated_qa]
+    query_embedding = embed_model.encode([query]).astype('float32')
+    curated_embeddings = embed_model.encode(curated_questions, show_progress_bar=False)
+    curated_embeddings = np.array(curated_embeddings).astype('float32')
+    
+    # Build a temporary FAISS index for curated questions
+    dimension = curated_embeddings.shape[1]
+    curated_index = faiss.IndexFlatL2(dimension)
+    curated_index.add(curated_embeddings)
+    
+    k = 1
+    distances, indices = curated_index.search(query_embedding, k)
+    
+    # Define a threshold for relevance (tune as needed)
+    threshold = 1.0
+    if distances[0][0] < threshold:
+        idx = indices[0][0]
+        qa_pair = curated_qa[idx]
+        return f"Curated Q/A Pair:\nQuestion: {qa_pair['question']}\nAnswer: {qa_pair['answer']}\n"
+    else:
+        return ""
+
 def main():
-    st.title("PDF Question-Answering App")
+    st.title("PDF Question-Answering App with Enhanced Context")
     
-    # Load FAISS index, chunks, and models
+    # Initialize conversation history if not already in session_state
+    if 'conversation_history' not in st.session_state:
+        st.session_state.conversation_history = ""
+    
+    # Load PDF index, chunks, and models
     index, chunks = load_index_and_chunks()
     embed_model = load_embedding_model()
     qa_pipeline = load_qa_pipeline()
@@ -39,29 +95,37 @@ def main():
     query = st.text_input("Question:")
     
     if query:
-        # Encode the query using the same SentenceTransformer model
-        query_embedding = embed_model.encode([query]).astype('float32')
+        # Append the current question to conversation history
+        st.session_state.conversation_history += f"User: {query}\n"
         
-        # Retrieve top k relevant chunks
-        k = 3
+        # Retrieve relevant PDF context using FAISS index
+        query_embedding = embed_model.encode([query]).astype('float32')
+        k = 3  # number of top chunks to retrieve
         distances, indices = index.search(query_embedding, k)
         
-        # Prepare combined context from the retrieved chunks
-        context = ""
+        pdf_context = ""
         for idx in indices[0]:
-            context_piece = chunks[idx]
-            context += context_piece + " "
+            pdf_context += chunks[idx] + "\n"
+        
+        # Get curated Q/A context if the query matches any curated questions
+        curated_context = get_curated_context(query, curated_qa_pairs, embed_model)
+        
+        # Combine conversation history, curated context, and PDF context for the QA pipeline
+        combined_context = st.session_state.conversation_history + "\n" + curated_context + "\n" + pdf_context
         
-        # Use an expander to optionally display the retrieved context
-        with st.expander("Show Retrieved Context"):
-            for idx in indices[0]:
-                st.write(chunks[idx])
+        # Optionally show the retrieved contexts in expanders
+        with st.expander("Show PDF Retrieved Context"):
+            st.write(pdf_context)
+        if curated_context:
+            with st.expander("Show Curated Q/A Context"):
+                st.write(curated_context)
         
         st.subheader("Answer:")
         try:
-            # Use the QA pipeline to generate an answer based on the combined context
-            result = qa_pipeline(question=query, context=context)
-            st.write(result["answer"])
+            result = qa_pipeline(question=query, context=combined_context)
+            answer = result["answer"]
+            st.write(answer)
+            st.session_state.conversation_history += f"AI: {answer}\n"
         except Exception as e:
             st.error(f"Error generating answer: {e}")