dangduytung commited on
Commit
d15ce79
1 Parent(s): 655b032

Add simple chatbot use model DiabloGPT

Browse files
Files changed (3) hide show
  1. __init__.py +4 -0
  2. main.py +82 -0
  3. requirements.txt +3 -0
__init__.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ MODEL_MICROSOFT_DIABLO_MEDIUM = 'microsoft/DialoGPT-medium'
2
+ MODEL_MICROSOFT_DIABLO_LARGE = 'microsoft/DialoGPT-large'
3
+
4
+ OUTPUT_MAX_LENGTH = 200
main.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import torch
4
+ import datetime
5
+ import __init__
6
+
7
+ MODEL_NAME = __init__.MODEL_MICROSOFT_DIABLO_MEDIUM
8
+ OUTPUT_MAX_LENGTH = __init__.OUTPUT_MAX_LENGTH
9
+
10
+
11
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
12
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
13
+
14
+
15
+ def print_f(session_id, text):
16
+ print(f"{datetime.datetime.now()} | {session_id} | {text}")
17
+
18
+
19
+ def predict(input, history, request: gr.Request):
20
+ session_id = 'UNKNOWN'
21
+ if request:
22
+ # Get session_id is client_ip + client_port
23
+ session_id = request.client.host + ':' + str(request.client.port)
24
+ # print_f(session_id, f" inp: {input}")
25
+
26
+ # Tokenize the new input sentence
27
+ new_user_input_ids = tokenizer.encode(
28
+ input + tokenizer.eos_token, return_tensors='pt')
29
+
30
+ # Append the new user input tokens to the chat history
31
+ bot_input_ids = torch.cat(
32
+ [torch.LongTensor(history), new_user_input_ids], dim=-1)
33
+
34
+ # Generate a response
35
+ history = model.generate(bot_input_ids, max_length=OUTPUT_MAX_LENGTH,
36
+ pad_token_id=tokenizer.eos_token_id).tolist()
37
+
38
+ # Convert the tokens to text, and then split the responses into lines
39
+ response = tokenizer.decode(history[0]).split("<|endoftext|>")
40
+
41
+ # Convert to tuples of list
42
+ response = [(response[i], response[i + 1])
43
+ for i in range(0, len(response) - 1, 2)]
44
+
45
+ # Print new conversation
46
+ print_f(session_id, response[-1])
47
+
48
+ return response, history
49
+
50
+
51
+ css = """
52
+ #row_bot{width: 70%; height: var(--size-96); margin: 0 auto}
53
+ #row_bot .block{background: var(--color-grey-100); height: 100%}
54
+ #row_input{width: 70%; margin: 0 auto}
55
+ #row_input .block{background: var(--color-grey-100)}
56
+
57
+ @media screen and (max-width: 768px) {
58
+ #row_bot{width: 100%; height: var(--size-96); margin: 0 auto}
59
+ #row_bot .block{background: var(--color-grey-100); height: 100%}
60
+ #row_input{width: 100%; margin: 0 auto}
61
+ #row_input .block{background: var(--color-grey-100)}
62
+ }
63
+ """
64
+ block = gr.Blocks(css=css, title="Chatbot")
65
+
66
+ with block:
67
+ gr.Markdown(f"""
68
+ <p style="font-size:20px; text-align: center">{MODEL_NAME}</p>
69
+ """)
70
+ with gr.Row(elem_id='row_bot'):
71
+ chatbot = gr.Chatbot()
72
+ with gr.Row(elem_id='row_input'):
73
+ message = gr.Textbox(placeholder="Enter something")
74
+ state = gr.State([])
75
+
76
+ message.submit(predict,
77
+ inputs=[message, state],
78
+ outputs=[chatbot, state])
79
+ message.submit(lambda x: "", message, message)
80
+
81
+ # Params ex: debug=True, share=True, server_name="0.0.0.0", server_port=5050
82
+ block.launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ transformers==4.22.2
2
+ torch==1.13.1
3
+ gradio==3.19.1