natdon/Michael_Scott_Bot

1from transformers import AutoTokenizer, AutoModelForCausalLM2import torch3import gradio as gr4 5tokenizer = AutoTokenizer.from_pretrained("natdon/DialoGPT_Michael_Scott")6model = AutoModelForCausalLM.from_pretrained("natdon/DialoGPT_Michael_Scott")7 8chat_history_ids = None9step = 010 11 12def predict(input, chat_history_ids=chat_history_ids, step=step):13    # encode the new user input, add the eos_token and return a tensor in Pytorch14    new_user_input_ids = tokenizer.encode(15        input + tokenizer.eos_token, return_tensors='pt')16 17    # append the new user input tokens to the chat history18    bot_input_ids = torch.cat(19        [chat_history_ids, new_user_input_ids], dim=-1) if step > 0 else new_user_input_ids20 21    # generated a response while limiting the total chat history to 1000 tokens,22    chat_history_ids = model.generate(23        bot_input_ids, max_length=1000,24        pad_token_id=tokenizer.eos_token_id,25        no_repeat_ngram_size=3,26        do_sample=True,27        top_k=100,28        top_p=0.7,29        temperature=0.830    )31    step = step + 132    output = tokenizer.decode(33        chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)34    return output35 36 37demo = gr.Blocks()38 39with demo:40    gr.Markdown(41        """42  <center>    43  <img src="https://media3.giphy.com/media/l0amJzVHIAfl7jMDos/giphy.gif" alt="dialog" width="250" height="250">44  45  ## Speak with Michael by typing in the input box below.46  </center>47  """48    )49 50    with gr.Row():51        with gr.Column():52            inp = gr.Textbox(53                label="Enter text to converse with Michael here:",54                lines=1,55                max_lines=1,56                value="Wow this is hard",57                placeholder="What do you think of Toby?",58            )59            btn = gr.Button("Submit")60        out = gr.Textbox(lines=3)61    # btn = gr.Button("Submit")62    inp.submit(fn=predict, inputs=inp, outputs=out)63    btn.click(fn=predict, inputs=inp, outputs=out)64 65demo.launch()66