Spaces:
Sleeping
Sleeping
import torch | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
model_path = "GRMenon/mental-mistral-7b-instruct-autotrain" | |
tokenizer = AutoTokenizer.from_pretrained(model_path) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_path, | |
device_map="auto", | |
torch_dtype='auto' | |
).eval() | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
# 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?"} | |
] | |
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: | |
print(response[0]) | |