Update src/streamlit_app.py

src/streamlit_app.py

@@ -11,201 +11,105 @@ logger = logging.getLogger(__name__)
 GROQ_API_KEY="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"
 OPENAI_API_KEY="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"
 
-logger.info("Design Page...")
-# -------------------------------
-# PAGE CONFIG (MUST BE FIRST)
-# -------------------------------
-PORT = int(os.environ.get("PORT", 8501))
-
-st.markdown("""
-<style>
-.main-title {
-    font-size: 52px;
-    font-weight: 800;
-    text-align: center;
-    color: #0B5ED7;
-    margin-bottom: 5px;
-}
-
-.sub-title {
-    font-size: 20px;
-    text-align: center;
-    color: #555555;
-    margin-bottom: 30px;
-}
-</style>
-""", unsafe_allow_html=True)
-
-st.markdown(
-    '<div class="main-title">💊 AI Medical Labelling System</div>',
-    unsafe_allow_html=True
-)
+ import os
+import logging
+from langchain_core.documents import Document
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from langchain_openai import ChatOpenAI
+import json
 
-st.markdown(
-    '<div class="sub-title">Simplifying FDA Drug Safety Information using Generative AI & RAG</div>',
-    unsafe_allow_html=True
+# Set up logging configuration
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
 )
 
-# -------------------------------
-# CUSTOM CSS (FANCY DESIGN)
-# -------------------------------
-st.markdown("""
-<style>
-.main {
-    background-color: #f7f9fc;
-}
-
-.big-title {
-    font-size:40px;
-    font-weight:700;
-    color:#1f4e79;
-}
-
-.subtitle {
-    font-size:18px;
-    color:#555;
-}
-
-.result-card {
-    background-color:white;
-    padding:20px;
-    border-radius:12px;
-    box-shadow:0px 2px 10px rgba(0,0,0,0.08);
-    margin-top:15px;
-}
-</style>
-""", unsafe_allow_html=True)
-
-# -------------------------------
-# HEADER
-# -------------------------------
-
-st.divider()
-
-# -------------------------------
-# SIDEBAR CONTROLS
-# -------------------------------
-with st.sidebar:
-    st.header("⚙️ Search Options")
-
-    drug_name = st.text_input(
-        "Drug Name",
-        placeholder="PHENYTOIN SODIUM"
-    )
-
-    selected_results = st.radio(
-        "Information Type",
-        ["Side Effects", "Warnings", "Both"]
-    )
-
-    run_button = st.button("🔍 Generate Explanation")
-
-# -------------------------------
-# LOAD ENV + MODELS
-# -------------------------------
-
-logger.info("Loading HuggingFace embedding model...")
-
-load_dotenv()
+# Get a logger for this module
+logger = logging.getLogger(__name__)
 
 working_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(working_dir)
+data_dir = f"{parent_dir}/data"
+vector_db_dir = f"{parent_dir}/vector_db"
 
-embeddings = HuggingFaceEmbeddings(
-    model_name="sentence-transformers/all-MiniLM-L6-v2"
-)
 
-vectordb = Chroma(
-    persist_directory=os.path.join(working_dir, "Chroma_db"),
-    embedding_function=embeddings
+logger.info("Reading Files Process Started...")
+all_records = []
+
+# loop through all files
+for file_name in os.listdir(data_dir):
+    if file_name.endswith(".json"):
+        file_path = os.path.join(data_dir, file_name)
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+            # if JSON contains list of records
+            if isinstance(data, list):
+                all_records.extend(data)
+            else:
+                all_records.append(data)
+
+print("Total drug records:", len(all_records))
+
+documents = []
+
+for record in data:
+
+    drug = record.get("generic_name", ["UNKNOWN"])[0].upper()
+
+    # choose sections you want in RAG
+    sections = [
+        "indications_and_usage",
+        "warnings_and_cautions",
+        "adverse_reactions",
+        "drug_interactions"
+    ]
+
+    for section in sections:
+        if section in record:
+
+            for text in record[section]:
+
+                documents.append(
+                    Document(
+                        page_content=text,
+                        metadata={
+                            "generic_name": drug,
+                            "section": section
+                        }
+                    )
+                )
+
+print("Documents created:", len(documents))
+
+logger.info("Split chunk Files Process Started...")
+splitter = RecursiveCharacterTextSplitter(
+    chunk_size=800,
+    chunk_overlap=150
 )
 
-logger.info("Calling OpenAI model gpt-4o-mini...")
+chunked_docs = splitter.split_documents(documents)
+
+print("Chunks created:", len(chunked_docs))
 
-llm = ChatOpenAI(
-    model="gpt-4o-mini",
-    temperature=0
+logger.info("Embeddings Files Process Started...")
+
+embeddings = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2"
+)
+#%%
+print("Chroma ready ✅")
+
+logger.info(" VectorDB Process Started...")
+vectordb = Chroma.from_documents(
+    documents=chunked_docs,
+    embedding=embeddings,
+    persist_directory="./chroma_db"
 )
 
-# -------------------------------
-# RAG FUNCTION
-# -------------------------------
-def generate_section(drug_name, section, rules):
-
-    results = vectordb.get(
-        where={
-            "$and": [
-                {"generic_name": drug_name},
-                {"section": section}
-            ]
-        }
-    )
-
-    documents = results.get("documents", [])
-
-    if not documents:
-        st.warning(f"No data found for {section}")
-        return
-
-    context = "\n".join(set(documents))
-
-    prompt = f"""
-You are a medical assistant.
-
-Rewrite the FDA drug information into simplified,
-easy-to-understand language.
-
-Rules:
-{rules}
-
-Drug: {drug_name}
-
-FDA TEXT:
-{context}
-"""
-
-    with st.spinner("🧠 AI is analysing FDA data..."):
-        response = llm.invoke(prompt)
-
-    st.markdown(
-        f'<div class="result-card">{response.content}</div>',
-        unsafe_allow_html=True
-    )
-
-logger.info("Configuring prompt..")
-# -------------------------------
-# RULES
-# -------------------------------
-SIDE_EFFECT_RULES = """
-- Use simple English
-- Bullet points (max 7)
-- Group similar side effects
-- Separate common vs serious
-"""
-
-WARNING_RULES = """
-- Use simple English
-- Bullet points (max 7)
-- Group warnings clearly
-"""
-
-SECTION_MAP = {
-    "Side Effects": [("adverse_reactions", SIDE_EFFECT_RULES)],
-    "Warnings": [("warnings_and_cautions", WARNING_RULES)],
-    "Both": [
-        ("adverse_reactions", SIDE_EFFECT_RULES),
-        ("warnings_and_cautions", WARNING_RULES),
-    ],
-}
-
-# -------------------------------
-# MAIN ACTION
-# -------------------------------
-if run_button and drug_name:
-
-    st.subheader(f"Results for: {drug_name.upper()}")
-
-    for section, rules in SECTION_MAP[selected_results]:
-        generate_section(drug_name, section, rules)
-
-elif run_button:
-    st.warning("Please enter a drug name.")
+print("Vector DB created successfully ✅")
+logger.info("VectorDB Process Completed...")
+