Update app.py

app.py

@@ -1,8 +1,6 @@
 # app.py
-# Synthetic Patient Records RAG App
+# Lightweight RAG App for Hugging Face Spaces (CPU-friendly)
 # Author: Your Name
-# Description: A Retrieval-Augmented Generation (RAG) application for synthetic hospital datasets.
-# Runs on Hugging Face Spaces using Gradio + FAISS + Sentence Transformers.
 
 import os
 import pandas as pd
@@ -15,18 +13,15 @@ from transformers import pipeline
 # ======================================================
 # 1. Dataset Handling
 # ======================================================
-
 DEFAULT_DATA_PATH = "patients.csv"
 
 def safe_load_csv(path):
-    """Safely load the dataset from CSV"""
     if not os.path.exists(path):
-        raise FileNotFoundError(f"❌ No dataset found at {path}. Please upload 'patients.csv'.")
+        raise FileNotFoundError(f"No dataset found at {path}. Please upload 'patients.csv'.")
     df = pd.read_csv(path)
     return df
 
 def preprocess_df(df):
-    """Cleans and harmonizes column names and fields"""
     df = df.copy()
     ren = {}
     for c in df.columns:
@@ -61,28 +56,16 @@ def preprocess_df(df):
         if col in df.columns:
             df[col] = pd.to_datetime(df[col], errors="coerce")
 
-    if "Length_of_Stay" not in df.columns or df["Length_of_Stay"].isnull().all():
-        if "Admission_Date" in df.columns and "Discharge_Date" in df.columns:
-            df["Length_of_Stay"] = (
-                (df["Discharge_Date"] - df["Admission_Date"]).dt.days.fillna(0).astype(int)
-            )
-        else:
-            df["Length_of_Stay"] = pd.Series([1] * len(df), index=df.index)
+    if "Length_of_Stay" not in df.columns:
+        df["Length_of_Stay"] = 1
+
+    diag_series = df["Diagnosis"].fillna("").astype(str) if "Diagnosis" in df.columns else pd.Series([""] * len(df))
+    treat_series = df["Treatment"].fillna("").astype(str) if "Treatment" in df.columns else pd.Series([""] * len(df))
 
-    diag_series = (
-        df["Diagnosis"].fillna("").astype(str) if "Diagnosis" in df.columns else pd.Series([""] * len(df))
-    )
-    treat_series = (
-        df["Treatment"].fillna("").astype(str) if "Treatment" in df.columns else pd.Series([""] * len(df))
-    )
     if "Notes" not in df.columns:
         df["Notes"] = (diag_series + " " + treat_series).str.strip()
 
     df["Notes"] = df["Notes"].astype(str)
-    df["Satisfaction_Score"] = pd.to_numeric(
-        df.get("Satisfaction_Score", pd.Series(np.nan, index=df.index)), errors="coerce"
-    ).fillna(-1)
-
     if "Patient_ID" not in df.columns:
         df.insert(0, "Patient_ID", range(1, len(df) + 1))
     return df.reset_index(drop=True)
@@ -91,10 +74,8 @@ def preprocess_df(df):
 # ======================================================
 # 2. Embedding + FAISS Setup
 # ======================================================
-
 def build_faiss_index(df, embed_model):
-    """Build FAISS index from Notes column"""
-    embeddings = embed_model.encode(df["Notes"].tolist(), convert_to_numpy=True, show_progress_bar=True)
+    embeddings = embed_model.encode(df["Notes"].tolist(), convert_to_numpy=True, show_progress_bar=False)
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(embeddings)
     return index, embeddings
@@ -103,16 +84,14 @@ def build_faiss_index(df, embed_model):
 # ======================================================
 # 3. RAG Query Function
 # ======================================================
-
 def generate_answer(query, df, embed_model, index, generator, top_k=3):
-    """Retrieve relevant notes and generate LLM summary"""
     query_emb = embed_model.encode([query])
     _, idxs = index.search(np.array(query_emb).astype("float32"), top_k)
     retrieved = df.iloc[idxs[0]]
     context = "\n".join(retrieved["Notes"].astype(str).tolist())
 
     prompt = f"""
-You are a hospital data assistant. Use the following patient notes to answer the user's question clearly and concisely.
+You are a hospital assistant. Use the following context to answer the question.
 
 Context:
 {context}
@@ -121,72 +100,59 @@ Question: {query}
 
 Answer:
 """
-    result = generator(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
+    result = generator(prompt, max_new_tokens=200)[0]["generated_text"]
     return result, retrieved[["Patient_ID", "Department", "Satisfaction_Score", "Length_of_Stay", "Notes"]]
 
 
 # ======================================================
-# 4. Gradio Interface Logic
+# 4. Gradio Interface
 # ======================================================
-
 def create_interface():
-    # Load default dataset
     df_raw = safe_load_csv(DEFAULT_DATA_PATH)
     df = preprocess_df(df_raw)
 
-    # Embeddings + Index
     embed_model = SentenceTransformer("all-MiniLM-L6-v2")
     index, _ = build_faiss_index(df, embed_model)
 
-    # LLM generator
-    generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3", device_map="auto")
+    # ✅ Use lightweight model (works on CPU)
+    generator = pipeline("text2text-generation", model="google/flan-t5-base")
 
     def query_app(user_query, uploaded_file=None):
-        """Handles user queries and optional dataset uploads"""
         local_df = df.copy()
         local_index = index
 
-        # If user uploads a dataset
         if uploaded_file is not None:
             try:
                 user_df = preprocess_df(pd.read_csv(uploaded_file.name))
                 local_df = user_df
                 local_index, _ = build_faiss_index(local_df, embed_model)
             except Exception as e:
-                return f"⚠️ Failed to process uploaded file: {str(e)}", pd.DataFrame()
+                return f"⚠️ Error loading file: {e}", pd.DataFrame()
 
-        # Run RAG pipeline
         answer, retrieved = generate_answer(user_query, local_df, embed_model, local_index, generator)
         return answer, retrieved
 
-    # Gradio UI
     iface = gr.Interface(
         fn=query_app,
         inputs=[
             gr.Textbox(label="💬 Ask a question about patient data"),
-            gr.File(label="📂 Upload a custom patient CSV (optional)")
+            gr.File(label="📂 Upload a patient CSV (optional)")
         ],
         outputs=[
             gr.Textbox(label="🤖 AI Generated Answer"),
             gr.Dataframe(label="📋 Retrieved Records")
         ],
         title="🏥 Synthetic Patient Records RAG App",
-        description=(
-            "Upload your patient dataset (or use the default one) and ask natural-language questions.\n"
-            "Built with Sentence Transformers + FAISS + Mistral 7B."
-        ),
+        description="Ask natural-language questions about synthetic hospital data. Powered by Sentence Transformers + Flan-T5.",
         examples=[
             ["Summarize satisfaction trends by department."],
-            ["Find patients over 65 with stays longer than 10 days."],
-            ["Generate a synthetic patient summary for a cardiology admission."],
-        ],
+            ["Find patients older than 65 with long hospital stays."],
+            ["Generate a summary of cardiology patients."]
+        ]
     )
     return iface
 
 
-# ======================================================
-# 5. Run App
-# ======================================================
 if __name__ == "__main__":
     app = create_interface()
     app.launch()