Update README.md

README.md

@@ -14,10 +14,10 @@ This model was trained using AutoTrain. For more information, please visit [Auto
 # Usage
 
 ```python
-
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_path = "PATH_TO_THIS_REPO"
+model_path = "GRMenon/Mental-Mistral-7B-Instruct-AutoTrain"
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 model = AutoModelForCausalLM.from_pretrained(
@@ -26,15 +26,17 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype='auto'
 ).eval()
 
-# Prompt content: "hi"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Prompt content:
 messages = [
-    {"role": "user", "content": "hi"}
+    {"role": "user", "content": "Hey Connor! I have been feeling a bit down lately. I could really use some advice on how to feel better?"}
 ]
 
-input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True, return_tensors='pt')
-output_ids = model.generate(input_ids.to('cuda'))
-response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+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)
 
-# Model response: "Hello! How can I assist you today?"
-print(response)
+# Model response: 
+print(response[0])
 ```