Views
No views yet
1from prompting.llama_prompt import modified_extes_support_strategies
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5def get_sys_msg_with_strategy(strategy):
6 if strategy is None:
7 return "You are a helpful, precise, and accurate emotional support expert."
8
9 description = modified_extes_support_strategies.get(strategy, "No description available")
10
11 return (f"You are a helpful and caring AI, which is an expert in emotional support. "
12 f"A user has come to you with emotional challenges, distress, or anxiety. "
13 f"Use the \"{strategy}\" strategy ({description}) for answering the user. "
14 "Make your response short and to the point.")
15
16cur_strategy = "Clarification"
17
18messages = [
19 {'role': 'system', 'content': get_sys_msg_with_strategy(cur_strategy)},
20 {'role': 'user', 'content': "Hello! How's it going?"},
21 {'role': 'assistant', 'content': 'Hello. How can I assist you today?'},
22 {'role': 'user', 'content': "I'm feeling a bit down today."},
23]
24
25device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
26
27model_name = "navidmadani/esconv_sra_llama3_8b"
28model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
29tokenizer = AutoTokenizer.from_pretrained(model_name)
30
31model = model.to(device)
32model.eval()
33
34input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
35input_t = torch.LongTensor([input_ids]).to(device)
36
37output = model.generate(input_t, max_new_tokens=512)[:, input_t.shape[1]:]
38resp = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
39print(resp)
40