Update app.py

app.py

@@ -1,38 +1,32 @@
-import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and tokenizer
-model_name = "GRMenon/mental-health-mistral-7b-instructv0.2-finetuned-V2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model_path = "GRMenon/mental-mistral-7b-instruct-autotrain"
 
-def get_bot_response(user_input):
-    # Tokenize user input
-    inputs = tokenizer.encode(user_input, return_tensors="pt")
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    device_map="auto",
+    torch_dtype='auto'
+).eval()
 
-    # Generate response
-    try:
-        bot_response = model.generate(inputs, max_length=100, num_return_sequences=1)
-    except Exception as e:
-        return f"Sorry, an error occurred while generating a response: {e}"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # Decode and return response
-    return tokenizer.decode(bot_response[0], skip_special_tokens=True)
+# Prompt content:
+messages = [
+    {"role": "user", "content": "Hey Connor! I have been feeling a bit down lately. I could really use some advice on how to feel better?"}
+]
 
-def main():
-    st.title("Mental Health Chatbot")
-    
-    st.write("Welcome to the mental health chatbot. You can ask questions or share your thoughts, and I'll try to provide helpful responses.")
-    
-    user_input = st.text_input("Type your message here:", value="", key="user_input")
-    
-    if user_input:
-        st.write(f"You: {user_input}")
-        
-        bot_response = get_bot_response(user_input)
-        
-        st.write(f"Chatbot: {bot_response}")
+input_ids = tokenizer.apply_chat_template(conversation=messages,
+                                          tokenize=True,
+                                          add_generation_prompt=True,
+                                          return_tensors='pt').to(device)
+output_ids = model.generate(input_ids=input_ids,
+                            max_new_tokens=512,
+                            do_sample=True,
+                            pad_token_id=2)
+response = tokenizer.batch_decode(output_ids.detach().cpu().numpy(),
+                                  skip_special_tokens = True)
 
-if __name__ == "__main__":
-    main()
+# Model response: 
+print(response[0])