26 lines
1.1 KiB
Python
26 lines
1.1 KiB
Python
|
from transformers.models.auto.tokenization_auto import AutoTokenizer
|
||
|
|
from transformers.models.auto.modeling_auto import AutoModelForCausalLM
|
||
|
|
import torch
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-large')
|
||
|
|
model = AutoModelForCausalLM.from_pretrained('microsoft/DialoGPT-large')
|
||
|
|
# model = AutoModelForCausalLM.from_pretrained('../output-medium')
|
||
|
|
|
||
|
|
# chatting 5 times with greedy search
|
||
|
|
for step in range(5):
|
||
|
|
# take user input
|
||
|
|
text = input(">> You: ")
|
||
|
|
# encode the input and add end of string token
|
||
|
|
input_ids = tokenizer.encode(text + tokenizer.eos_token, return_tensors="pt")
|
||
|
|
# concatenate new user input with chat history (if there is)
|
||
|
|
bot_input_ids = torch.cat([chat_history_ids, input_ids], dim=-1) if step > 0 else input_ids
|
||
|
|
# generate a bot response
|
||
|
|
chat_history_ids = model.generate(
|
||
|
|
bot_input_ids,
|
||
|
|
max_length=1000,
|
||
|
|
pad_token_id=tokenizer.eos_token_id,
|
||
|
|
)
|
||
|
|
#print the output
|
||
|
|
output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
|
||
|
|
print(f"Bot: {output}")
|