Update app.py

app.py

@@ -1,5 +1,14 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, pipeline
+import streamlit as st
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+    pipeline,
+)
+#from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, pipeline
 #from llama_cpp import Llama
 from datasets import load_dataset
 import os
@@ -47,24 +56,32 @@ st.markdown('<div class="blurred-background"></div>', unsafe_allow_html=True)
 
 #"""""""""""""""""""""""""   Application Code Starts here   """""""""""""""""""""""""""""""""""""""""""""
 
-# Load the dataset
+# Cache resource for dataset loading
 @st.cache_resource
 def load_counseling_dataset():
-    return load_dataset("Amod/mental_health_counseling_conversations")
+    # Load a smaller subset of the dataset for memory efficiency
+    dataset = load_dataset("Amod/mental_health_counseling_conversations", split="train")
+    return dataset
 
-dataset = load_counseling_dataset()
+# Process the dataset in batches to avoid memory overuse
+def process_dataset_in_batches(dataset, batch_size=1000):
+    for example in dataset.shuffle().select(range(batch_size)):
+        yield example
 
 # Fine-tune the model and save it
 @st.cache_resource
 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)
 
+    # Enable gradient checkpointing for memory optimization
+    model.gradient_checkpointing_enable()
+
     # Prepare dataset for training
+    dataset = load_counseling_dataset()
+    
     def preprocess_function(examples):
         return tokenizer(examples["context"] + "\n" + examples["response"], truncation=True)
 
@@ -74,13 +91,16 @@ def fine_tune_model():
     # Training arguments
     training_args = TrainingArguments(
         output_dir="./fine_tuned_model",
-        evaluation_strategy="epoch",
+        evaluation_strategy="steps",
         learning_rate=2e-5,
-        per_device_train_batch_size=1,
+        per_device_train_batch_size=8,
+        per_device_eval_batch_size=8,
         num_train_epochs=3,
         weight_decay=0.01,
+        fp16=True,  # Enable FP16 for lower memory usage
         save_total_limit=2,
-        save_strategy="epoch"
+        save_steps=500,
+        logging_steps=100,
     )
 
     # Trainer
@@ -124,10 +144,8 @@ if st.button("Get Supportive Response"):
     if user_input.strip():
         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"]
+                response = pipe(user_input, max_length=150, num_return_sequences=1)[0]["generated_text"]
                 st.subheader("Supportive Suggestion:")
                 st.markdown(f"**{response}**")
             except Exception as e: