Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 #from transformers import AutoTokenizer
 from llama_cpp import Llama
-#from transformers import pipeline
+from transformers import pipeline
 #from peft import PeftModel, PeftConfig
 #from transformers import AutoModelForCausalLM
 from datasets import load_dataset
@@ -45,24 +45,20 @@ st.markdown(
 # Add the blurred background div
 st.markdown('<div class="blurred-background"></div>', unsafe_allow_html=True)
 
-"""""""""""""""""""""""""   Application Code Starts here   """""""""""""""""""""""""""""""""""""""""""""
+#"""""""""""""""""""""""""   Application Code Starts here   """""""""""""""""""""""""""""""""""""""""""""
 
-# Path to the GGUF model on Hugging Face (Hugging Face Spaces automatically downloads it)
-MODEL_PATH = "/root/.cache/huggingface/hub/models--QuantFactory--Mental-Health-FineTuned-Mistral-7B-Instruct-v0.2-GGUF/blobs/"
-
-# Load Llama model
+# Load the text generation model pipeline
 @st.cache_resource
-def load_llama_model():
+def load_text_generation_model():
     try:
-        # The GGUF model will be cached automatically by llama.cpp
-        return Llama(model_path=f"{MODEL_PATH}/model.gguf", n_threads=8)  # Adjust n_threads based on your environment
+        return pipeline("text-generation", model="QuantFactory/Mental-Health-FineTuned-Mistral-7B-Instruct-v0.2-GGUF")
     except Exception as e:
         st.error(f"Error loading model: {e}")
         return None
 
-llama_model = load_llama_model()
+text_generator = load_text_generation_model()
 
-# Load dataset for context
+# Load the counseling dataset
 @st.cache_resource
 def load_counseling_dataset():
     return load_dataset("Amod/mental_health_counseling_conversations")
@@ -77,8 +73,8 @@ This platform is designed to provide **supportive, positive, and encouraging res
 """)
 
 # Check if the model loaded correctly
-if llama_model is None:
-    st.error("The text generation model could not be loaded. Please check your configuration.")
+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"):
@@ -94,12 +90,14 @@ else:
     if st.button("Get Supportive Response"):
         if user_input.strip():
             try:
-                # Generate response using Llama
+                # Generate response using the text generation pipeline
                 prompt = f"User: {user_input}\nCounselor:"
-                response = llama_model(prompt, max_tokens=200, stop=["\n", "User:"])
+                response = text_generator(prompt, max_length=200, num_return_sequences=1)
                 
+                # Extract and display the response
+                counselor_reply = response[0]["generated_text"].split("Counselor:")[-1].strip()
                 st.subheader("Counselor's Response:")
-                st.write(response["choices"][0]["text"].strip())
+                st.write(counselor_reply)
             except Exception as e:
                 st.error(f"An error occurred while generating the response: {e}")
         else: