fix multi-round dialogue.

generate.py

@@ -72,7 +72,9 @@ def __init__(self, model, tokenizer):
  self.model = model
  self.tokenizer = tokenizer
 
- def generate(self, args, prompts):
+ def generate(self, args, prompts, cut_off=None, cut_off_times=1):
+ if cut_off is not None:
+ cut_off_times = [cut_off_times for i in range(len(prompts))]
  batch = len(prompts)
  assert batch <= args.batch_size
 
@@ -118,6 +120,12 @@ def generate(self, args, prompts):
  try:
  t.index(self.tokenizer.eos_id)
  except ValueError:
+ if cut_off is not None:
+ if cut_off == self.tokenizer.decode(t[:cur_pos + 1])[-len(cut_off):]:
+ if cut_off_times[i] == 1:
+ continue
+ else:
+ cut_off_times[i] -= 1
  continue_exsample.append(i)
  if len(continue_exsample) == 0:
  break
@@ -126,6 +134,7 @@ def generate(self, args, prompts):
  for i, t in enumerate(tokens.tolist()):
  t = t[: args.seq_length]
  try:
+ t = t[: t.index(self.tokenizer.pad_id)]
  t = t[: t.index(self.tokenizer.eos_id)]
  except ValueError:
  pass

llama_dialogue.py

@@ -21,15 +21,15 @@ def multi_round_chat(args, lm_generation, keep_length_ratio=0.5):
 
  input_str = ''
  for user, ans in zip(users, answers):
- input_str += user + '\n' + ans + '\n'
- input_str += user_input + '\n'
+ input_str += 'User: ' + user + '\nBot: ' + ans + '\n'
+ input_str += 'User: ' + user_input + '\nBot: '
  if len(input_str) >= int(keep_length_ratio * args.seq_length):
  input_str = input_str[len(input_str) - int(keep_length_ratio * args.seq_length):]
- answer = lm_generation.generate(args, [input_str])[0]
+ answer = lm_generation.generate(args, [input_str], cut_off='User:', cut_off_times=1)[0]
  answer = answer[len(input_str):]
- print("ChatLLaMa: " + answer + '\n')
- users.append(user_input)
- answers.append(answer)
+ print("ChatLLaMa: " + answer.replace('User:', ''))
+ users.append(user_input.rstrip(' ').rstrip('\n'))
+ answers.append(answer.replace('User:', '').rstrip(' ').rstrip('\n'))
 
 
 if __name__ == '__main__':