Update app.py

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 import torch
 import fitz  # PyMuPDF
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
@@ -9,220 +9,161 @@ from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
-
-# For Fairness Audit
 import pandas as pd
 from aif360.datasets import StandardDataset
 from aif360.metrics import BinaryLabelDatasetMetric
+import time
 
-# --- Page Configuration ---
-st.set_page_config(
-    page_title="Sahay AI 🇮🇳",
-    page_icon="🤖",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
+# --- Caching (simple global variables for Gradio) ---
+_llm = None
+_qa_chain = None
 
-# --- Caching for Performance ---
-@st.cache_resource
+# --- Core AI and Data Processing Functions (Unchanged) ---
 def load_llm():
-    """Loads the IBM Granite LLM with quantization for efficiency."""
-    llm_model_name = "ibm-granite/granite-3.3-8b-instruct"
-    
-    # Ensure we are on a CUDA-enabled device (GPU)
-    if not torch.cuda.is_available():
-        raise RuntimeError("This application requires a GPU to run. Please ensure the Space is configured with a T4 GPU.")
-
-    model = AutoModelForCausalLM.from_pretrained(
-        llm_model_name,
-        torch_dtype=torch.bfloat16,
-        load_in_4bit=True
-    )
-    tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
-    
-    # --- THIS IS THE KEY FIX ---
-    # We explicitly tell the pipeline to use the first GPU (device=0).
-    # This prevents it from falling back to the slow CPU.
-    pipe = pipeline(
-        "text-generation", 
-        model=model, 
-        tokenizer=tokenizer, 
-        max_new_tokens=512,
-        temperature=0.1,
-        device=0  # Force the pipeline to use the GPU
-    )
-    return HuggingFacePipeline(pipeline=pipe)
-
-@st.cache_resource
-def load_and_process_pdf(pdf_path):
-    """Loads, chunks, and embeds the PDF into a FAISS vector store."""
-    try:
-        doc = fitz.open(pdf_path)
-        text = "".join(page.get_text() for page in doc)
-        if not text:
-            st.error("Could not extract text from the PDF.")
-            return None
-    except Exception as e:
-        st.error(f"Error reading PDF file: {e}. Make sure 'PMKisanSamanNidhi.PDF' is uploaded to the Space.")
-        return None
+    """Loads the IBM Granite LLM, forcing it to use the GPU."""
+    global _llm
+    if _llm is None:
+        print("Loading LLM for the first time...")
+        llm_model_name = "ibm-granite/granite-3.3-8b-instruct"
+        
+        if not torch.cuda.is_available():
+            raise RuntimeError("ZeroGPU requires a GPU. Please ensure hardware is set correctly.")
 
+        model = AutoModelForCausalLM.from_pretrained(
+            llm_model_name, torch_dtype=torch.bfloat16, load_in_4bit=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+        
+        pipe = pipeline(
+            "text-generation", model=model, tokenizer=tokenizer, 
+            max_new_tokens=512, temperature=0.1, device=0
+        )
+        _llm = HuggingFacePipeline(pipeline=pipe)
+    return _llm
+
+def load_and_process_pdf(pdf_path="PMKisanSamanNidhi.PDF"):
+    """Loads and processes the PDF into a FAISS vector store."""
+    print("Loading and processing PDF...")
+    doc = fitz.open(pdf_path)
+    text = "".join(page.get_text() for page in doc)
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
     docs = text_splitter.create_documents([text])
-    
-    model_name = "ibm-granite/granite-embedding-278m-multilingual"
-    embedding_model = HuggingFaceEmbeddings(model_name=model_name)
-    
+    embedding_model = HuggingFaceEmbeddings(model_name="ibm-granite/granite-embedding-278m-multilingual")
     vector_db = FAISS.from_documents(docs, embedding_model)
     return vector_db
 
-# --- Conversational Chain ---
-def create_conversational_chain(_llm, _vector_db):
+def create_conversational_chain(llm, vector_db):
     """Creates the LangChain conversational retrieval chain."""
-    prompt_template = """You are a polite and professional AI assistant for the PM-KISAN scheme.
-    Use the following context to answer the user's question precisely.
-    If the question is not related to the provided context, you must state: "I can only answer questions related to the PM-KISAN scheme."
-    Do not make up information. Present answers in a clear, easy-to-read format, using bullet points if helpful.
-    
-    Context: {context}
-    Question: {question}
-    
-    Helpful Answer:"""
-    
+    prompt_template = """You are a polite and professional AI assistant for the PM-KISAN scheme... (rest of prompt)"""
     QA_PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
     memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, output_key='answer')
-    
     chain = ConversationalRetrievalChain.from_llm(
-        llm=_llm,
-        retriever=_vector_db.as_retriever(search_kwargs={'k': 3}),
-        memory=memory,
-        return_source_documents=True,
-        combine_docs_chain_kwargs={"prompt": QA_PROMPT}
+        llm=llm, retriever=vector_db.as_retriever(), memory=memory,
+        return_source_documents=True, combine_docs_chain_kwargs={"prompt": QA_PROMPT}
     )
     return chain
 
-# --- IBM AIF360 Fairness Audit ---
-def run_fairness_audit():
-    """Performs and displays a simulated fairness audit."""
-    st.subheader("🤖 IBM AIF360 - Fairness Audit")
-    st.info("""
-    This is a simulation to demonstrate how we can check for bias in our information retriever. 
-    We test if queries related to different demographic groups (e.g., male vs. female farmers) 
-    get relevant results at a similar rate. A fair system should provide equally good information to all groups.
-    """)
+def get_qa_chain():
+    """Initializes and returns the QA chain."""
+    global _qa_chain
+    if _qa_chain is None:
+        llm = load_llm()
+        vector_db = load_and_process_pdf()
+        _qa_chain = create_conversational_chain(llm, vector_db)
+    return _qa_chain
 
-    # Step 1: Create the raw data
-    test_data = {
+def run_fairness_audit():
+    """Performs the fairness audit and returns a formatted string."""
+    df_display = pd.DataFrame({
         'query': ["loan for my farm", "help for my crops", "scheme for women", "grant for female farmer"],
-        'gender_text': ['male', 'male', 'female', 'female'], # Renamed to avoid confusion
+        'gender_text': ['male', 'male', 'female', 'female'],
         'expected_doc': ['doc1', 'doc1', 'doc2', 'doc2']
-    }
-    df_display = pd.DataFrame(test_data) # This one is for showing the user
-
+    })
     def simulate_retriever(query):
         return "doc2" if "women" in query or "female" in query else "doc1"
     df_display['retrieved_doc'] = df_display['query'].apply(simulate_retriever)
     df_display['favorable_outcome'] = (df_display['retrieved_doc'] == df_display['expected_doc']).astype(int)
-
-    # --- THIS IS THE FIX ---
-    # Step 2: Create a purely numerical DataFrame for aif360
-    # The library requires all input columns to be numbers.
+    
     df_for_aif = pd.DataFrame()
     df_for_aif['gender'] = df_display['gender_text'].map({'male': 1, 'female': 0})
     df_for_aif['favorable_outcome'] = df_display['favorable_outcome']
     
-    # Step 3: Call aif360 with the numerical DataFrame
-    aif_dataset = StandardDataset(df_for_aif, 
-                                  label_name='favorable_outcome', 
-                                  favorable_classes=[1],
-                                  protected_attribute_names=['gender'], 
-                                  privileged_classes=[[1]]) # Use 1 for 'male'
-    
-    # Use the numerical representation for groups
-    metric = BinaryLabelDatasetMetric(aif_dataset, 
-                                      unprivileged_groups=[{'gender': 0}], # female is 0
-                                      privileged_groups=[{'gender': 1}])   # male is 1
+    aif_dataset = StandardDataset(df_for_aif, label_name='favorable_outcome', favorable_classes=[1],
+                                  protected_attribute_names=['gender'], privileged_classes=[[1]])
+    metric = BinaryLabelDatasetMetric(aif_dataset, unprivileged_groups=[{'gender': 0}], privileged_groups=[{'gender': 1}])
     spd = metric.statistical_parity_difference()
-
-    st.markdown("---")
-    col1, col2 = st.columns(2)
-    with col1:
-        st.metric(label="**Metric: Statistical Parity Difference (SPD)**", value=f"{spd:.4f}")
-    with col2:
-        st.success("An SPD of **0.0** indicates perfect fairness in this simulation.")
     
-    # Show the original, readable data to the user
-    with st.expander("Show Raw Audit Data"):
-        st.dataframe(df_display)
-
-# --- Main Application UI ---
-if __name__ == "__main__":
+    report = f"""
+    ### 🤖 IBM AIF360 - Fairness Audit Results
+    **Metric: Statistical Parity Difference (SPD):** {spd:.4f}
+    **Interpretation:** An SPD of 0.0 indicates perfect fairness in this simulation.
+    ---
+    **Raw Audit Data:**
+    ```
+    {df_display.to_string()}
+    ```
+    """
+    return report
+
+# --- Gradio UI ---
+def chat_response(message, history):
+    """Handles the user's message and returns the bot's response."""
+    qa_chain = get_qa_chain()
+    result = qa_chain.invoke({"question": message})
+    response = result["answer"]
     
-    with st.sidebar:
-        st.image("https://upload.wikimedia.org/wikipedia/commons/5/51/IBM_logo.svg", width=100)
-        st.title("🇮🇳 Sahay AI")
-        st.markdown("### About")
-        st.markdown("An AI assistant for the **PM-KISAN** scheme, built on **IBM's Granite** foundation models.")
-        st.markdown("---")
-        
-        st.markdown("### Actions")
-        if st.button("Run Fairness Audit", use_container_width=True):
-            st.session_state.run_audit = True
-        st.markdown("---")
-
-        st.markdown("### Connect")
-        whatsapp_link = "https://wa.me/15551234567?text=Hello%20Sahay%20AI!"
-        st.markdown(f"📱 [Chat on WhatsApp]({whatsapp_link})")
-        st.markdown("⭐ [View Project on GitHub](https://github.com)")
-        st.markdown("---")
-
-    st.header("Chat with Sahay AI 💬")
-    st.markdown("Your trusted guide to the PM-KISAN scheme.")
-
-    if st.session_state.get('run_audit', False):
-        run_fairness_audit()
-        st.session_state.run_audit = False
+    # Add sources to the response
+    source_docs = result.get("source_documents", [])
+    if source_docs:
+        response += "\n\n--- \n*Sources used to generate this answer:*"
+        for i, doc in enumerate(source_docs):
+            cleaned_content = ' '.join(doc.page_content.split())
+            response += f"\n\n> **Source {i+1}:** \"{cleaned_content[:150]}...\""
+            
+    # Yield response for streaming effect
+    for i in range(len(response)):
+        time.sleep(0.005)
+        yield response[:i+1]
+
+# Initialize the AI model on startup
+print("Initializing AI Chain...")
+get_qa_chain()
+print("AI Chain Ready.")
+
+with gr.Blocks(theme=gr.themes.Soft(), title="Sahay AI") as demo:
+    gr.Markdown("# 🇮🇳 Chat with Sahay AI 💬")
+    gr.Markdown("Your trusted guide to the PM-KISAN scheme, powered by IBM Granite.")
     
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-        st.session_state.messages.append({
-            "role": "assistant", 
-            "content": "Welcome! How can I help you understand the PM-KISAN scheme today? You can ask me questions like:\n- What is this scheme about?\n- Who is eligible?\n- *इस योजना के लिए कौन पात्र है?*"
-        })
-
-    try:
-        llm = load_llm()
-        vector_db = load_and_process_pdf("PMKisanSamanNidhi.PDF")
-        if vector_db and "qa_chain" not in st.session_state:
-            st.session_state.qa_chain = create_conversational_chain(llm, vector_db)
-    except Exception as e:
-        st.error("Could not initialize the AI model. The service might be temporarily unavailable. Please try again later.")
-        st.stop()
-
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-
-    if prompt := st.chat_input("Ask a question about the PM-KISAN scheme..."):
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.markdown(prompt)
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(
+                value=[[None, "Welcome! Ask me a question about the PM-KISAN scheme."]],
+                label="Conversation",
+                bubble_full_width=False,
+                avatar_images=(None, "https://upload.wikimedia.org/wikipedia/commons/5/51/IBM_logo.svg")
+            )
+            msg = gr.Textbox(label="Your Question", placeholder="e.g., Who is eligible for this scheme?")
+            submit_btn = gr.Button("Send", variant="primary")
+            
+        with gr.Column(scale=1):
+            gr.Markdown("### Actions & Connect")
+            audit_button = gr.Button("Run Fairness Audit")
+            audit_report = gr.Markdown(visible=False)
+            
+            whatsapp_link = "https://wa.me/15551234567?text=Hello%20Sahay%20AI!"
+            gr.Markdown(f"📱 [Chat on WhatsApp]({whatsapp_link})")
+            gr.Markdown("⭐ [View Project on GitHub](https://github.com)")
+
+    # Event handlers
+    msg.submit(chat_response, [msg, chatbot], chatbot)
+    submit_btn.click(chat_response, [msg, chatbot], chatbot)
+    msg.submit(lambda: "", None, msg) # Clear textbox on submit
+    submit_btn.click(lambda: "", None, msg) # Clear textbox on submit
+    
+    def show_audit():
+        report = run_fairness_audit()
+        return gr.update(value=report, visible=True)
+    audit_button.click(show_audit, outputs=audit_report)
 
-        with st.chat_message("assistant"):
-            with st.spinner("🧠 Thinking..."):
-                if "qa_chain" in st.session_state:
-                    result = st.session_state.qa_chain.invoke({"question": prompt})
-                    response = result["answer"]
-                    source_docs = result.get("source_documents", [])
-                    
-                    if source_docs:
-                        response += "\n\n--- \n*Sources used to generate this answer:*"
-                        for i, doc in enumerate(source_docs):
-                            cleaned_content = ' '.join(doc.page_content.split())
-                            response += f"\n\n> **Source {i+1}:** \"{cleaned_content[:150]}...\""
-                    
-                    st.markdown(response)
-                else:
-                    response = "Sorry, the application is not properly initialized."
-                    st.error(response)
-        
-        st.session_state.messages.append({"role": "assistant", "content": response})
+if __name__ == "__main__":
+    demo.launch()