Update app.py

app.py

@@ -47,79 +47,99 @@ st.markdown('<div class="blurred-background"></div>', unsafe_allow_html=True)
 
 #"""""""""""""""""""""""""   Application Code Starts here   """""""""""""""""""""""""""""""""""""""""""""
 
-# Groq API Configuration
-api_key = os.environ.get("LawersGuideAPIKey")  # Ensure GROQ_API_KEY is set in your environment variables
-base_url = "https://api.groq.com/openai/v1/models/google/gemma-2-9b-it/completions"
+# Load the dataset
+@st.cache_resource
+def load_counseling_dataset():
+    return load_dataset("Amod/mental_health_counseling_conversations")
 
-headers = {
-    "Authorization": f"Bearer {api_key}",
-    "Content-Type": "application/json"
-}
+dataset = load_counseling_dataset()
 
-# Function to query Groq model
+# Fine-tune the model and save it
 @st.cache_resource
-def query_groq_model(prompt, max_tokens=100, temperature=0.7):
-    try:
-        payload = {
-            "prompt": prompt,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-            "top_p": 1.0,
-            "frequency_penalty": 0.0,
-            "presence_penalty": 0.0,
-            "n": 1
-        }
-        response = requests.post(base_url, headers=headers, json=payload)
-        response.raise_for_status()
-        result = response.json()
-        return result["choices"][0]["text"].strip()
-    except Exception as e:
-        return f"Error querying the model: {e}"
+def fine_tune_model():
+    from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling
+    
+    # Load base model and tokenizer
+    model_name = "prabureddy/Mental-Health-FineTuned-Mistral-7B-Instruct-v0.2"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Streamlit App
-st.title("Mental Health Counseling Chat")
-st.markdown("""
-Welcome to the **Mental Health Counseling Chat Application**.  
-This platform is designed to provide **supportive, positive, and encouraging responses** using the Groq `google/gemma-2-9b-it` model.
-""")
+    # Prepare dataset for training
+    def preprocess_function(examples):
+        return tokenizer(examples["context"] + "\n" + examples["response"], truncation=True)
+
+    tokenized_datasets = dataset.map(preprocess_function, batched=True)
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir="./fine_tuned_model",
+        evaluation_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=1,
+        num_train_epochs=3,
+        weight_decay=0.01,
+        save_total_limit=2,
+        save_strategy="epoch"
+    )
+
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["validation"],
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
 
-# Load example dataset for user exploration (optional)
+    trainer.train()
+
+    # Save the fine-tuned model
+    trainer.save_model("./fine_tuned_model")
+    tokenizer.save_pretrained("./fine_tuned_model")
+    return "./fine_tuned_model"
+
+# Load or fine-tune the model
+model_dir = fine_tune_model()
+
+# Load the fine-tuned model for inference
 @st.cache_resource
-def load_counseling_dataset():
-    from datasets import load_dataset
-    return load_dataset("Amod/mental_health_counseling_conversations")
+def load_pipeline(model_dir):
+    return pipeline("text-generation", model=model_dir)
 
-dataset = load_counseling_dataset()
+pipe = load_pipeline(model_dir)
 
-# Display example questions and answers from dataset
-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("---")
+# Streamlit App
+st.title("Mental Health Support Assistant")
+st.markdown("""
+Welcome to the **Mental Health Support Assistant**.  
+This tool helps detect potential mental health concerns based on user input and provides **uplifting and positive suggestions** to boost morale.
+""")
 
 # User input for mental health concerns
-user_input = st.text_area("Your question or concern:", placeholder="Type your question here...")
+user_input = st.text_area("Please share your concern:", placeholder="Type your question or concern here...")
 
 if st.button("Get Supportive Response"):
     if user_input.strip():
-        try:
-            # Query Groq model
-            prompt = f"User: {user_input}\nCounselor:"
-            counselor_reply = query_groq_model(prompt, max_tokens=150, temperature=0.7)
-            st.subheader("Counselor's Response:")
-            st.write(counselor_reply)
-        except Exception as e:
-            st.error(f"An error occurred while querying the model: {e}")
+        with st.spinner("Analyzing your input and generating a response..."):
+            try:
+                # Construct the messages for the pipeline
+                messages = [{"role": "user", "content": user_input}]
+                # Generate a response
+                response = pipe(messages)[0]["generated_text"]
+                st.subheader("Supportive Suggestion:")
+                st.markdown(f"**{response}**")
+            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.")
+        st.error("Please enter a concern to receive suggestions.")
 
-# Sidebar resources
-st.sidebar.header("Additional Mental Health Resources")
+# Sidebar for additional resources
+st.sidebar.header("Additional 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.")
+st.sidebar.info("This application is not a replacement for professional counseling. If you're in crisis, seek professional help immediately.")