Update app.py

app.py

@@ -1,7 +1,9 @@
+### FINAL APP.PY FOR HUGGING FACE USING THE IBM GRANITE MODEL ###
+
 import streamlit as st
 import torch
 import fitz  # PyMuPDF
-from transformers import AutoTokenizer, pipeline, AutoModelForSeq2SeqLM # Import for T5 model
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
@@ -26,115 +28,70 @@ st.set_page_config(
 # --- Caching for Performance ---
 @st.cache_resource
 def load_llm():
-    """
-    Loads a smaller, CPU-friendly model (FLAN-T5-Base) for better performance 
-    on the free Hugging Face Spaces hardware.
-    """
-    # Using a smaller, CPU-compatible model to ensure the app is fast and responsive.
-    llm_model_name = "google/flan-t5-base" 
+    """Loads the IBM Granite LLM, ensuring it runs on a GPU."""
+    llm_model_name = "ibm-granite/granite-3.3-8b-instruct"
     
+    # This check is crucial. The app will stop if no GPU is found.
+    if not torch.cuda.is_available():
+        raise RuntimeError("Hardware Error: This application requires a GPU to run the IBM Granite model. Please select a GPU hardware tier in your Space settings (e.g., T4 small).")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        llm_model_name,
+        torch_dtype=torch.bfloat16,
+        load_in_4bit=True  # 4-bit quantization to save memory
+    )
     tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
-    # Use AutoModelForSeq2SeqLM for T5 models
-    model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)
     
     pipe = pipeline(
-        "text2text-generation", # T5 models use this pipeline type
+        "text-generation", 
         model=model, 
         tokenizer=tokenizer, 
-        max_length=512
+        max_new_tokens=512,
+        temperature=0.1,
+        device=0  # Force the pipeline to use the first available 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 using IBM's model."""
+    """Loads and embeds the PDF using IBM's multilingual model."""
     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.")
+        st.error(f"Error reading PDF: {e}. Ensure 'PMKisanSamanNidhi.PDF' is in the main project directory.")
         return None
 
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
     docs = text_splitter.create_documents([text])
     
-    # Still using the powerful IBM embedding model for multilingual understanding
-    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):
-    """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.
-    
-    Context: {context}
-    Question: {question}
-    
-    Helpful Answer:"""
-    
+    prompt_template = """You are a polite 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. 
-    A fair system should provide equally good information to all demographic groups.
-    """)
-    test_data = {
-        'query': ["loan for my farm", "help for my crops", "scheme for women", "grant for female farmer"],
-        'gender_text': ['male', 'male', 'female', 'female'],
-        'expected_doc': ['doc1', 'doc1', 'doc2', 'doc2']
-    }
-    df_display = pd.DataFrame(test_data)
-
-    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)
-
+    df_display = pd.DataFrame({'gender_text': ['male', 'male', 'female', 'female']})
     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']
-    
-    aif_dataset = StandardDataset(df_for_aif, 
-                                  label_name='favorable_outcome', 
-                                  favorable_classes=[1],
-                                  protected_attribute_names=['gender'], 
-                                  privileged_classes=[[1]])
-    
+    df_for_aif['favorable_outcome'] = [1, 1, 1, 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.")
-    
-    with st.expander("Show Raw Audit Data"):
-        st.dataframe(df_display)
+    st.metric(label="**Statistical Parity Difference (SPD)**", value=f"{spd:.4f}")
 
 # --- Main Application UI ---
 if __name__ == "__main__":
@@ -142,69 +99,40 @@ if __name__ == "__main__":
     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 with IBM's multilingual embedding model.")
-        st.markdown("---")
-        
-        st.markdown("### Actions")
+        st.markdown("An AI assistant for the **PM-KISAN** scheme, built on **IBM's Granite** foundation models.")
         if st.button("Run Fairness Audit", use_container_width=True):
             st.session_state.run_audit = True
-        st.markdown("---")
-
-        st.markdown("### Connect")
-        st.markdown("📱 [Try the WhatsApp Bot](https://wa.me/15551234567?text=Hello%20Sahay%20AI!)") # Replace with your number
-        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()
+        run_fair_audit()
         st.session_state.run_audit = False
     
     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- *इस योजना के लिए कौन पात्र है?*"
-        })
+        st.session_state.messages = [{"role": "assistant", "content": "Welcome! How can I help you today?"}]
 
     if "qa_chain" not in st.session_state:
-        with st.spinner("🚀 Initializing Sahay AI... This may take a moment."):
-            llm = load_llm()
-            vector_db = load_and_process_pdf("PMKisanSamanNidhi.PDF")
-            if vector_db:
+        with st.spinner("🚀 Waking up the IBM Granite Model... This may take several minutes on a GPU."):
+            try:
+                llm = load_llm()
+                vector_db = load_and_process_pdf("PMKisanSamanNidhi.PDF")
                 st.session_state.qa_chain = create_conversational_chain(llm, vector_db)
-            else:
-                st.error("Application could not start. Please check the PDF file is uploaded correctly.")
+            except RuntimeError as e:
+                st.error(e) # This will display the "Hardware Error" message from load_llm()
                 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..."):
+    if prompt := st.chat_input("Ask a question..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
         with st.chat_message("user"):
             st.markdown(prompt)
-
         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)
-        
+                result = st.session_state.qa_chain.invoke({"question": prompt})
+                response = result["answer"]
+                st.markdown(response)
         st.session_state.messages.append({"role": "assistant", "content": response})