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]) | |