shumingma commited on
Commit
6280e9d
·
1 Parent(s): dc4f215

update readme

Browse files
Files changed (1) hide show
  1. README.md +4 -3
README.md CHANGED
@@ -92,11 +92,12 @@ messages = [
92
  {"role": "system", "content": "You are a helpful AI assistant."},
93
  {"role": "user", "content": "How are you?"},
94
  ]
95
- chat_input = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
 
96
 
97
  # Generate response
98
- chat_outputs = model.generate(chat_input, max_new_tokens=50)
99
- response = tokenizer.decode(chat_outputs[0][chat_input.shape[-1]:], skip_special_tokens=True) # Decode only the response part
100
  print("\nAssistant Response:", response)
101
  ```
102
 
 
92
  {"role": "system", "content": "You are a helpful AI assistant."},
93
  {"role": "user", "content": "How are you?"},
94
  ]
95
+ prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
96
+ chat_input = tokenizer(prompt, return_tensors="pt").to(model.device)
97
 
98
  # Generate response
99
+ chat_outputs = model.generate(**chat_input, max_new_tokens=50)
100
+ response = tokenizer.decode(chat_outputs[0][chat_input['input_ids'].shape[-1]:], skip_special_tokens=True) # Decode only the response part
101
  print("\nAssistant Response:", response)
102
  ```
103