Create app.py

app.py

@@ -0,0 +1,274 @@
+# ======================================================================================
+# 1. SETUP: Patch SQLite and Import Libraries
+# ======================================================================================
+# This MUST be the first import to ensure ChromaDB uses the correct SQLite version
+import sys
+import os
+os.environ['PYSQLITE3_BUNDLED'] = '1'
+__import__('pysqlite3')
+sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
+
+# Standard and third-party libraries
+import json
+import pandas as pd
+from typing import List, Union
+
+import chromadb
+
+import gradio as gr
+from pydantic import BaseModel, ValidationError
+from sentence_transformers import SentenceTransformer, CrossEncoder
+
+# LangChain imports
+from langchain_openai.chat_models import ChatOpenAI
+from langchain_community.vectorstores import Chroma
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema.output_parser import StrOutputParser
+from langchain.output_parsers import PydanticOutputParser
+from langchain_community.embeddings import SentenceTransformerEmbeddings
+
+# ======================================================================================
+# 2. CONSTANTS AND CONFIGURATION
+# ======================================================================================
+DB_DIR = "./chroma_db"
+COLLECTION_NAME = "clinical_examples"
+EMBEDDING_MODEL_NAME = "pritamdeka/S-Biomed-Roberta-snli-multinli-stsb"
+RERANKER_MODEL_NAME = 'cross-encoder/ms-marco-MiniLM-L-6-v2'
+DATASET_URL = "https://huggingface.co/datasets/DanFed/patient_encounters1_notes_preprocessed/raw/main/patient_encounters1_notes_preprocessed.csv"
+
+
+# ======================================================================================
+# 3. DATABASE SETUP: One-time data loading and embedding
+# ======================================================================================
+def setup_database(client: chromadb.Client):
+    """
+    Loads data, generates embeddings, and populates the ChromaDB collection
+    only if it's empty.
+    """
+    collection = client.get_or_create_collection(name=COLLECTION_NAME)
+
+    if collection.count() > 0:
+        print(f"Collection '{COLLECTION_NAME}' already exists with {collection.count()} documents. Skipping setup.")
+        return
+
+    print(f"Collection '{COLLECTION_NAME}' is empty. Starting data population...")
+
+    # Load dataset
+    df = pd.read_csv(DATASET_URL)
+    df.drop(['index', 'ENCOUNTER_ID', 'CLINICAL_NOTES', 'BIRTHDATE', 'FIRST',
+           'START', 'STOP', 'PATIENT_ID', 'ENCOUNTERCLASS', 'CODE', 'DESCRIPTION',
+           'BASE_ENCOUNTER_COST', 'TOTAL_CLAIM_COST', 'PAYER_COVERAGE',
+           'REASONCODE', 'REASONDESCRIPTION', 'PATIENT_AGE',
+           'DESCRIPTION_OBSERVATIONS', 'DESCRIPTION_CONDITIONS',
+           'DESCRIPTION_MEDICATIONS', 'DESCRIPTION_PROCEDURES', 'AGE_GROUP'], axis=1, inplace=True)
+
+    # Create example strings
+    def create_examples(row):
+        return f"Message: \n\n{row['ENCOUNTER_PROMPT'].strip()}\n\nResult: \n\n{row['COND_MED_PRO_STRUCTURED'].strip()}"
+    df['EXAMPLES'] = df.apply(create_examples, axis=1)
+
+    # Generate embeddings
+    model = SentenceTransformer(EMBEDDING_MODEL_NAME)
+    examples = df["EXAMPLES"].tolist()
+    embeddings = model.encode(
+        examples,
+        batch_size=32,
+        show_progress_bar=True,
+        convert_to_numpy=True
+    )
+
+    # Add to collection
+    collection.add(
+        documents=df["EXAMPLES"].tolist(),
+        embeddings=embeddings.tolist(),
+        ids=[str(i) for i in range(len(df["EXAMPLES"]))]
+    )
+    print(f"Successfully added {len(df['EXAMPLES'])} documents to the '{COLLECTION_NAME}' collection.")
+
+
+# ======================================================================================
+# 4. APPLICATION GLOBALS AND AI COMPONENTS
+# ======================================================================================
+# Pydantic schema for structured output
+class ClinicalExtraction(BaseModel):
+    conditions: List[str]
+    medications: List[str]
+    procedures: List[str]
+
+# Parser and format instructions
+parser = PydanticOutputParser(pydantic_object=ClinicalExtraction)
+format_instructions = parser.get_format_instructions().replace("{", "{{").replace("}", "}}")
+
+# Global variables for AI components
+LANGCHAIN_LLM = None
+FINAL_PROMPT = None
+FINAL_CHAIN = None
+VECTOR_STORE = None
+RERANKER = CrossEncoder(RERANKER_MODEL_NAME)
+
+def initialize_ai_components(api_key: str):
+    """Initializes all AI components needed for the RAG pipeline."""
+    global LANGCHAIN_LLM, FINAL_PROMPT, FINAL_CHAIN
+    if not api_key:
+        raise gr.Error("OpenAI API Key is required!")
+
+    # LLM
+    LANGCHAIN_LLM = ChatOpenAI(openai_api_key=api_key, temperature=0.2)
+
+    # Prompt Template
+    FINAL_PROMPT = ChatPromptTemplate.from_template(
+        f"""You are a clinical information extractor.
+Extract EXACTLY this JSON format and nothing else:
+
+{format_instructions}
+
+CONTEXT (examples):
+
+{{context}}
+
+INPUT MESSAGE (clinical note + surrounding metadata):
+
+{{input}}
+
+Result:"""
+    )
+
+    # RAG Chain
+    FINAL_CHAIN = (
+        {"context": RunnablePassthrough(), "input": RunnablePassthrough()}
+        | FINAL_PROMPT
+        | LANGCHAIN_LLM
+        | StrOutputParser()
+    )
+    return "<p style='color:green;'>AI components initialized successfully!</p>"
+
+# ======================================================================================
+# 5. RAG PIPELINE
+# ======================================================================================
+def format_docs(docs):
+    """Join doc.page_content with blank lines."""
+    return "\n\n".join(d.page_content for d in docs)
+
+def generate_rag_response(input_text: str) -> Union[dict, str]:
+    """
+    Performs retrieval, reranking, generation, and validation.
+    """
+    if not FINAL_CHAIN or not VECTOR_STORE:
+        return "Error: AI components not initialized. Please set your API key."
+
+    # Initial embedding retrieval (top 20)
+    retriever = VECTOR_STORE.as_retriever(search_kwargs={"k": 20})
+    candidates = retriever.get_relevant_documents(input_text)
+
+    # Cross-encoder rerank -> top 5
+    pairs  = [(input_text, d.page_content) for d in candidates]
+    scores = RERANKER.predict(pairs)
+    sorted_docs = [d for _, d in sorted(zip(scores, candidates), reverse=True)]
+    top_docs = sorted_docs[:5]
+
+    # Build context and invoke chain
+    context = format_docs(top_docs)
+    raw_output = FINAL_CHAIN.invoke({"context": context, "input": input_text})
+
+    # Parse and validate the output
+    try:
+        parsed = parser.parse(raw_output)
+        return parsed.dict()
+    except ValidationError as e:
+        return f"Schema validation failed: {e}. Raw output was: {raw_output}"
+
+# ======================================================================================
+# 6. GRADIO UI
+# ======================================================================================
+def create_gradio_ui():
+    """Defines and returns the Gradio UI blocks."""
+    with gr.Blocks(title="Clinical Information Extractor") as demo:
+        gr.Markdown("# Clinical Information Extractor with RAG and Reranking")
+
+        with gr.Accordion("API Key Configuration", open=True):
+            key_box = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")
+            key_btn = gr.Button("Set API Key")
+            key_status = gr.Markdown("")
+        key_btn.click(initialize_ai_components, inputs=[key_box], outputs=[key_status])
+
+        gr.Markdown("---")
+        gr.Markdown("## Enter Clinical Note and Metadata")
+
+        with gr.Row():
+            age_group_input = gr.Textbox(label="Age Group", placeholder="e.g., middle adulthood")
+            visit_type_input = gr.Textbox(label="Visit Type", placeholder="e.g., ambulatory")
+        description_input = gr.Textbox(label="Description", placeholder="e.g., encounter for check up (procedure)")
+        note_input = gr.Textbox(label="Clinical Note", placeholder="Type the clinical note here...", lines=5)
+        
+        chatbot = gr.Chatbot(label="Extraction History", height=400)
+        send_btn = gr.Button("➡️ Extract Information")
+
+        def chat_interface(age, visit, desc, note, history):
+            history = history or []
+            
+            # Build full input with metadata
+            metadata_parts = []
+            if age: metadata_parts.append(f"Age group: {age}")
+            if visit: metadata_parts.append(f"Visit type: {visit}")
+            if desc: metadata_parts.append(f"Description: {desc}")
+            metadata_str = " | ".join(metadata_parts)
+            
+            full_input = f"{metadata_str}\n\nClinical Note:\n{note}" if metadata_str else note
+            user_display = f"**Metadata**: {metadata_str}\n\n**Note**: {note}"
+            
+            # Get response from RAG pipeline
+            response = generate_rag_response(full_input)
+            
+            # Format bot response
+            if isinstance(response, dict):
+                bot_response = f"```json\n{json.dumps(response, indent=2)}\n```"
+            else:
+                bot_response = str(response)
+            
+            history.append((user_display, bot_response))
+            return history, "" # Return updated history and clear the input textbox
+
+        send_btn.click(
+            fn=chat_interface,
+            inputs=[age_group_input, visit_type_input, description_input, note_input, chatbot],
+            outputs=[chatbot, note_input]
+        )
+        note_input.submit(
+            fn=chat_interface,
+            inputs=[age_group_input, visit_type_input, description_input, note_input, chatbot],
+            outputs=[chatbot, note_input]
+        )
+    return demo
+
+# ======================================================================================
+# 7. MAIN EXECUTION
+# ======================================================================================
+def main():
+    """
+    Main function to set up the database, initialize components, and launch the UI.
+    """
+    global VECTOR_STORE
+    
+    # 1. Setup ChromaDB client
+    client = chromadb.PersistentClient(path=DB_DIR)
+    
+    # 2. Populate the database if needed
+    setup_database(client)
+    
+    # 3. Initialize the LangChain vector store wrapper
+    embeddings = SentenceTransformerEmbeddings(model_name=EMBEDDING_MODEL_NAME)
+    VECTOR_STORE = Chroma(
+        client=client,
+        collection_name=COLLECTION_NAME,
+        embedding_function=embeddings,
+    )
+    print(f"Vector store initialized with {VECTOR_STORE._collection.count()} documents.")
+    
+    # 4. Create and launch the Gradio UI
+    demo = create_gradio_ui()
+    print("Launching Clinical IE Demo...")
+    demo.launch(server_name="0.0.0.0")
+
+if __name__ == "__main__":
+    main()