Spaces:
Sleeping
Sleeping
File size: 3,526 Bytes
3aa655b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import re
model_name = 'armandnlp/gpt2-TOD_finetuned_SGD'
tokenizer_TOD = AutoTokenizer.from_pretrained(model_name)
model_TOD = AutoModelForCausalLM.from_pretrained(model_name)
def generate_response(prompt):
input_ids = tokenizer_TOD(prompt, return_tensors="pt").input_ids
outputs = model_TOD.generate(input_ids,
do_sample=False,
max_length=1024,
eos_token_id=50262)
return tokenizer_TOD.batch_decode(outputs)[0]
#<|context|> <|user|> I want to go to the restaurant.<|endofcontext|>
import gradio as gr
iface = gr.Interface(fn=generate_response,
inputs="text",
outputs="text",
title="gpt2-TOD",
examples=[["<|context|> <|user|> I'm super hungry ! I want to go to the restaurant.<|endofcontext|>"],
["<|context|> <|user|> I want to go to the restaurant. <|system|> What food would you like to eat ? <|user|> Italian sounds good. <|endofcontext|>"]],
description="Passing in a task-oriented dialogue context generates a belief state, actions to take and a response based on those actions",
)
iface.launch()
"""
## Work in progress
## https://gradio.app/creating_a_chatbot/
## make chatbot interface
## can get input and responses for now
## would like to add belief state and actions to history response
## means modifying the history when appending input during next turn
## ie. keeping only the response and adding <|system|> token
ckpt = 'armandnlp/gpt2-TOD_finetuned_SGD'
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt)
def predict(input, history=[]):
# history: list of all token ids
# response: list of tuples of strings corresponding to dialogue turns
#model input and output with extra formatting
new_user_input_ids = tokenizer.encode(' <|user|> '+input, return_tensors='pt')
context = tokenizer.encode('<|context|>', return_tensors='pt')
endofcontext = tokenizer.encode(' <|endofcontext|>', return_tensors='pt')
model_input = torch.cat([context, torch.LongTensor(history), new_user_input_ids, endofcontext], dim=-1)
out = model.generate(model_input, max_length=1024, eos_token_id=50262).tolist()[0]
#history : format for next dialogue turn
history = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
string_out = tokenizer.decode(out)
response_only = string_out.split('<|response|>')[1].replace('<|endofresponse|>', '')
resp_tokenized = tokenizer.encode(' <|system|> '+response_only, return_tensors='pt')
history = torch.cat([history, resp_tokenized], dim=-1).tolist()
# history with belief + action
# output with belief + action + response
#response: format printed output
turns = tokenizer.decode(history[0])
#turns = "<|user|> I want to go to the restaurant. <|system|> What food would you like to eat ? <|user|> Italian sounds good. <|system|> Okay then !"
turns = re.split('<\|system\|>|<\|user\|>', turns)[1:]
#print(turns)
response = [(turns[i], turns[i+1]) for i in range(0, len(turns)-1, 2)]
#print(response)
return response, history
#predict("I want to go to the restaurant.")
import gradio as gr
gr.Interface(fn=predict,
inputs=["text", "state"],
outputs=["chatbot", "state"]).launch()
"""
|