File size: 1,583 Bytes
3ef14f8
 
caeef30
3ef14f8
 
 
a233482
d498112
3ef14f8
 
 
 
 
 
a5971a0
a233482
 
 
 
 
 
 
 
 
 
a5971a0
3ef14f8
 
 
a233482
 
 
 
 
 
 
 
 
 
 
 
30aaf10
 
d506c2a
30aaf10
86a2758
30aaf10
 
a233482
a5971a0
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
from transformers import AutoTokenizer, AutoModelForCausalLM

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]

def chat(message, history):
    history = history or []

    
    output = generate_response(message)
    context, response = output.split('<|endofcontext|>')

    history.append((context+'<|endofcontext|>', response))
    
    return history


import gradio as gr

chatbot = gr.Chatbot(color_map=("gray", "blue"))

iface = gr.Interface(chat,
                    ["text", "state"],
                    [chatbot, "state"],
                    allow_screenshot=False,
                    allow_flagging="never",
)



"""
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|>"]],
                     description="Passing in a task-oriented dialogue context generates a belief state, actions to take and a response based on those actions",
                     )
"""
iface.launch()