DialogGPT Bot

bot.py

import torch

i = 0
maxlen = 1024
while True:
    user_input = input('>> User: ').strip()
    
    if user_input.lower() == "q":
        break
    
    input_ids = tokenizer(user_input + tokenizer.eos_token, return_tensors='pt').input_ids
    input_ids = torch.cat([response, input_ids], dim=-1) if i > 0 else input_ids
    
    response = model.generate(input_ids, 
                              max_length=maxlen,
                              pad_token_id=tokenizer.eos_token_id)

    # As in the GPT-2 example above, the generated text includes the "prompt", so
    # we remove it. Then we decode as above.
    input_len = input_ids.shape[-1]
    generated = tokenizer.decode(response[:, input_len:][0], skip_special_tokens=True)
    print("DialoGPT: {}".format(generated))
    i += 1