Update app.py

app.py

@@ -1,9 +1,7 @@
 import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 #from transformers import AutoTokenizer
 #from llama_cpp import Llama
-from transformers import pipeline
-#from peft import PeftModel, PeftConfig
-#from transformers import AutoModelForCausalLM
 from datasets import load_dataset
 
 # Replace with the direct image URL
@@ -48,22 +46,22 @@ st.markdown('<div class="blurred-background"></div>', unsafe_allow_html=True)
 #"""""""""""""""""""""""""   Application Code Starts here   """""""""""""""""""""""""""""""""""""""""""""
 
 # Hugging Face access token
-HF_TOKEN = "HF_TOKEN"  # Replace with your actual token or use environment variables
+HF_TOKEN = "HF_TOKEN"  # Replace with your actual token or set it as an environment variable
 
-# Load the text generation model pipeline with authentication
+# Load the text generation pipeline with model and tokenizer
 @st.cache_resource
-def load_text_generation_model():
-    try:
-        return pipeline(
-            "text-generation",
-            model="google/gemma-2-9b-it",
-            use_auth_token=HF_TOKEN
-        )
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        return None
+def load_text_generation_pipeline():
+    model_name = "google/gemma-2-9b-it"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        load_in_8bit=True, 
+        device_map="auto", 
+        use_auth_token=HF_TOKEN
+    )
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-text_generator = load_text_generation_model()
+text_generator = load_text_generation_pipeline()
 
 # Load the counseling dataset
 @st.cache_resource
@@ -79,43 +77,38 @@ Welcome to the **Mental Health Counseling Chat Application**.
 This platform is designed to provide **supportive, positive, and encouraging responses** using a fast and efficient language model.
 """)
 
-# Check if the model loaded correctly
-if text_generator is None:
-    st.error("The text generation model could not be loaded. Please check your Hugging Face configuration.")
-else:
-    # Explore dataset for additional context or resources (optional)
-    if st.checkbox("Show Example Questions and Answers from Dataset"):
-        sample = dataset["train"].shuffle(seed=42).select(range(3))  # Display 3 random samples
-        for example in sample:
-            st.markdown(f"**Question:** {example.get('context', 'N/A')}")
-            st.markdown(f"**Answer:** {example.get('response', 'N/A')}")
-            st.markdown("---")
+# Display example dataset entries
+if st.checkbox("Show Example Questions and Answers from Dataset"):
+    sample = dataset["train"].shuffle(seed=42).select(range(3))  # Display 3 random samples
+    for example in sample:
+        st.markdown(f"**Question:** {example.get('context', 'N/A')}")
+        st.markdown(f"**Answer:** {example.get('response', 'N/A')}")
+        st.markdown("---")
 
-    # User input for mental health concerns
-    user_input = st.text_area("Your question or concern:", placeholder="Type your question here...")
+# User input for mental health concerns
+user_input = st.text_area("Your question or concern:", placeholder="Type your question here...")
 
-    if st.button("Get Supportive Response"):
-        if user_input.strip():
-            try:
-                # Generate response using the text generation pipeline
-                prompt = f"User: {user_input}\nCounselor:"
-                response = text_generator(prompt, max_length=200, num_return_sequences=1)
-                
-                # Extract and display the response
-                counselor_reply = response[0]["generated_text"].strip()
-                st.subheader("Counselor's Response:")
-                st.write(counselor_reply)
-            except Exception as e:
-                st.error(f"An error occurred while generating the response: {e}")
-        else:
-            st.error("Please enter a question or concern to receive a response.")
+if st.button("Get Supportive Response"):
+    if user_input.strip():
+        try:
+            # Generate response using the text generation pipeline
+            prompt = f"User: {user_input}\nCounselor:"
+            response = text_generator(prompt, max_length=100, num_return_sequences=1)
+            
+            # Extract and display the response
+            counselor_reply = response[0]["generated_text"].strip()
+            st.subheader("Counselor's Response:")
+            st.write(counselor_reply)
+        except Exception as e:
+            st.error(f"An error occurred while generating the response: {e}")
+    else:
+        st.error("Please enter a question or concern to receive a response.")
 
-    # Sidebar resources
-    st.sidebar.header("Additional Mental Health Resources")
-    st.sidebar.markdown("""
-    - [Mental Health Foundation](https://www.mentalhealth.org)
-    - [Mind](https://www.mind.org.uk)
-    - [National Suicide Prevention Lifeline](https://suicidepreventionlifeline.org)
-    """)
-
-    st.sidebar.info("This application is not a replacement for professional counseling. If you are in crisis, please seek professional help immediately.")
+# Sidebar resources
+st.sidebar.header("Additional Mental Health Resources")
+st.sidebar.markdown("""
+- [Mental Health Foundation](https://www.mentalhealth.org)
+- [Mind](https://www.mind.org.uk)
+- [National Suicide Prevention Lifeline](https://suicidepreventionlifeline.org)
+""")
+st.sidebar.info("This application is not a replacement for professional counseling. If you are in crisis, please seek professional help immediately.")