File size: 10,387 Bytes
74b8229
 
 
 
 
 
 
b4c0306
 
 
74b8229
 
766eaec
 
74b8229
766eaec
 
 
74b8229
766eaec
 
b4c0306
766eaec
 
74b8229
 
766eaec
74b8229
766eaec
74b8229
 
 
 
 
0d868fb
74b8229
766eaec
74b8229
766eaec
74b8229
 
 
 
 
 
 
a5629da
74b8229
 
f6e8a16
a5629da
 
 
 
 
f6e8a16
74b8229
235585a
74b8229
235585a
 
 
 
b4c0306
235585a
74b8229
 
38ca40a
814a362
38ca40a
 
 
 
b4c0306
38ca40a
814a362
74b8229
 
1173b30
74b8229
 
 
 
 
 
 
 
 
 
 
a5629da
2318c3b
74b8229
6e1a316
2318c3b
b4c0306
 
235585a
74b8229
235585a
 
 
 
 
 
a5629da
235585a
 
 
 
b4c0306
235585a
74b8229
 
 
 
0b3d061
b4c0306
 
0b3d061
b4c0306
0b3d061
 
b4c0306
74b8229
 
 
 
 
 
 
 
0b3d061
 
a5629da
74b8229
 
 
 
 
98bcb65
a5629da
74b8229
22887e0
74b8229
 
 
a5629da
74b8229
 
 
 
 
 
 
 
 
 
 
 
 
 
 
86d9965
74b8229
 
0d868fb
a5629da
0d868fb
 
 
896d60b
0d868fb
 
 
74b8229
 
 
a5629da
74b8229
a5629da
74b8229
 
 
 
 
 
 
 
b8c7507
 
 
 
a5629da
b8c7507
 
74b8229
 
 
 
 
 
b8c7507
 
 
 
 
74b8229
 
0d868fb
74b8229
b8c7507
a5629da
74b8229
 
 
 
 
 
 
 
 
 
a5629da
74b8229
d080a18
ab8d5cf
74b8229
a5629da
74b8229
 
38ca40a
766eaec
 
 
b4c0306
766eaec
 
38ca40a
896d60b
38ca40a
 
146762d
a5629da
 
 
 
 
38ca40a
74b8229
bc8f2a7
74b8229
b4c0306
 
 
 
 
 
 
 
 
 
74b8229
3fdc337
 
74b8229
814a362
 
 
b4c0306
 
74b8229
814a362
74b8229
 
 
 
 
 
 
 
 
 
b4c0306
74b8229
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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
"""
app.py - the main file for the app. This creates the flask app and handles the routes.

"""

import argparse
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")

import os
import sys
import time
import warnings
from os.path import dirname
from pathlib import Path

import gradio as gr
import nltk
import torch
from cleantext import clean
from gradio.inputs import Slider, Textbox, Radio
from transformers import pipeline

from converse import discussion
from grammar_improve import (
    build_symspell_obj,
    detect_propers,
    fix_punct_spacing,
    load_ns_checker,
    neuspell_correct,
    remove_repeated_words,
    remove_trailing_punctuation,
    symspeller,
    synthesize_grammar,
)
from utils import corr

nltk.download("stopwords")  # download stopwords

sys.path.append(dirname(dirname(os.path.abspath(__file__))))
warnings.filterwarnings(action="ignore", message=".*gradient_checkpointing*")
import transformers

transformers.logging.set_verbosity_error()
cwd = Path.cwd()
_cwd_str = str(cwd.resolve())  # string so it can be passed to os.path() objects


def chat(
    prompt_message,
    temperature: float = 0.5,
    top_p: float = 0.95,
    top_k: int = 20,
    constrained_generation: str = "False",
) -> str:
    """
    chat - the main function for the chatbot. This is the function that is called when the user

    :param _type_ prompt_message: the message to send to the model
    :param float temperature: the temperature value for the model, defaults to 0.6
    :param float top_p: the top_p value for the model, defaults to 0.95
    :param int top_k: the top_k value for the model, defaults to 25
    :param bool constrained_generation: whether to use constrained generation or not, defaults to False
    :return str: the response from the model
    """
    history = []
    response = ask_gpt(
        message=prompt_message,
        chat_pipe=my_chatbot,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        constrained_generation="true" in constrained_generation.lower(),
    )
    history = [prompt_message, response]
    html = ""
    for item in history:
        html += f"<b>{item}</b> <br><br>"

    html += ""

    return html


def ask_gpt(
    message: str,
    chat_pipe,
    speaker="person alpha",
    responder="person beta",
    min_length=12,
    max_length=48,
    top_p=0.95,
    top_k=25,
    temperature=0.5,
    constrained_generation=False,
    max_input_length=128,
) -> str:
    """
    ask_gpt - helper function that asks the GPT model a question and returns the response

    :param str message: the question to ask the model
    :param chat_pipe: the pipeline object for the model, created by the pipeline() function
    :param str speaker: the name of the speaker, defaults to "person alpha"
    :param str responder: the name of the responder, defaults to "person beta"
    :param int min_length: the minimum length of the response, defaults to 12
    :param int max_length: the maximum length of the response, defaults to 64
    :param float top_p: the top_p value for the model, defaults to 0.95
    :param int top_k: the top_k value for the model, defaults to 25
    :param float temperature: the temperature value for the model, defaults to 0.6
    :param bool constrained_generation: whether to use constrained generation or not, defaults to False
    :return str: the response from the model
    """
    st = time.perf_counter()
    prompt = clean(message)  # clean user input
    prompt = prompt.strip()  # get rid of any extra whitespace
    in_len = len(chat_pipe.tokenizer(prompt).input_ids)
    if in_len > max_input_length:
        # truncate to last max_input_length tokens
        tokens = chat_pipe.tokenizer(prompt).input_ids
        trunc_tokens = tokens[-max_input_length:]
        prompt = chat_pipe.tokenizer.decode(trunc_tokens)
        print(f"truncated prompt to {len(trunc_tokens)} tokens, input length: {in_len}")
    logging.info(f"prompt: {prompt}")
    resp = discussion(
        prompt_text=prompt,
        pipeline=chat_pipe,
        speaker=speaker,
        responder=responder,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        max_length=max_length,
        min_length=min_length,
        constrained_beam_search=constrained_generation,
    )
    gpt_et = time.perf_counter()
    gpt_rt = round(gpt_et - st, 2)
    rawtxt = resp["out_text"]
    # check for proper nouns
    if basic_sc:
        cln_resp = symspeller(rawtxt, sym_checker=basic_spell)
    else:
        cln_resp = synthesize_grammar(corrector=grammarbot, message=rawtxt)
    bot_resp_a = corr(remove_repeated_words(cln_resp))
    bot_resp = fix_punct_spacing(bot_resp_a)
    corr_rt = round(time.perf_counter() - gpt_et, 4)
    print(f"{gpt_rt + corr_rt} to respond, {gpt_rt} GPT, {corr_rt} for correction\n")
    return remove_trailing_punctuation(bot_resp)


def get_parser():
    """
    get_parser - a helper function for the argparse module
    """
    parser = argparse.ArgumentParser(
        description="submit a question, GPT model responds"
    )
    parser.add_argument(
        "-m",
        "--model",
        required=False,
        type=str,
        default="ethzanalytics/ai-msgbot-gpt2-XL",  # default model
        help="the model to use for the chatbot on https://huggingface.co/models OR a path to a local model",
    )
    parser.add_argument(
        "-gm",
        "--gram-model",
        required=False,
        type=str,
        default="pszemraj/grammar-synthesis-base",
        help="text2text generation model ID from huggingface for the model to correct grammar",
    )

    parser.add_argument(
        "--basic-sc",
        required=False,
        default=False,
        action="store_true",
        help="use symspell (statistical spelling correction) instead of neural spell correction",
    )

    parser.add_argument(
        "--verbose",
        action="store_true",
        default=False,
        help="turn on verbose logging",
    )
    parser.add_argument(
        "--test",
        action="store_true",
        default=False,
        help="load the smallest model for simple testing (ethzanalytics/distilgpt2-tiny-conversational)",
    )

    return parser


if __name__ == "__main__":
    args = get_parser().parse_args()
    default_model = str(args.model)
    test = args.test
    if test:
        logging.info("loading the smallest model for testing")
        default_model = "ethzanalytics/distilgpt2-tiny-conversational"

    model_loc = Path(default_model)  # if the model is a path, use it
    basic_sc = args.basic_sc  # whether to use the baseline spellchecker
    gram_model = str(args.gram_model)
    device = 0 if torch.cuda.is_available() else -1

    logging.info(f"CUDA avail is {torch.cuda.is_available()}")

    my_chatbot = (
        pipeline("text-generation", model=model_loc.resolve(), device=device)
        if model_loc.exists() and model_loc.is_dir()
        else pipeline("text-generation", model=default_model, device=device)
    )  # if the model is a name, use it. stays on CPU if no GPU available
    print(f"using model {my_chatbot.model}")

    if basic_sc:
        print("Using the baseline spellchecker")
        basic_spell = build_symspell_obj()
    else:
        print("using neural spell checker")
        grammarbot = pipeline("text2text-generation", gram_model, device=device)

    logging.info(f"using model stored here: \n {model_loc} \n")
    iface = gr.Interface(
        chat,
        inputs=[
            Textbox(
                default="Why is everyone here eating chocolate cake?",
                label="prompt_message",
                placeholder="Start a conversation with the bot",
                lines=2,
            ),
            Slider(
                minimum=0.0, maximum=1.0, step=0.05, default=0.4, label="temperature"
            ),
            Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.95, label="top_p"),
            Slider(minimum=0, maximum=100, step=5, default=20, label="top_k"),
            Radio(
                choices=["True", "False"],
                default="False",
                label="constrained_generation",
            ),
        ],
        outputs="html",
        examples_per_page=8,
        examples=[
            ["Point Break or Bad Boys II?", 0.75, 0.95, 50, False],
            ["So... you're saying this wasn't an accident?", 0.6, 0.95, 40, False],
            ["Hi, my name is Reginald", 0.6, 0.95, 100, False],
            ["Happy birthday!", 0.9, 0.95, 50, False],
            ["I have a question, can you help me?", 0.6, 0.95, 50, False],
            ["Do you know a joke?", 0.8, 0.85, 50, False],
            ["Will you marry me?", 0.9, 0.95, 100, False],
            ["Are you single?", 0.95, 0.95, 100, False],
            ["Do you like people?", 0.7, 0.95, 25, False],
            ["You never took a shortcut before?", 0.7, 0.95, 100, False],
        ],
        title=f"GPT Chatbot Demo: {default_model} Model",
        description=f"A Demo of a Chatbot trained for conversation with humans. Size XL= 1.5B parameters.\n\n"
        "**Important Notes & About:**\n\n"
        "You can find a link to the model card **[here](https://huggingface.co/ethzanalytics/ai-msgbot-gpt2-XL-dialogue)**\n\n"
        "1. responses can take up to 60 seconds to respond sometimes, patience is a virtue.\n"
        "2. the model was trained on several different datasets.  fact-check responses instead of regarding as a true statement.\n"
        "3. Try adjusting the **[generation parameters](https://huggingface.co/blog/how-to-generate)** to get a better understanding of how they work!\n"
        "4. New - try using [constrained beam search](https://huggingface.co/blog/constrained-beam-search) decoding to generate more coherent responses. _(experimental, feedback welcome!)_\n",
        css="""
            .chatbox {display:flex;flex-direction:row}
            .user_msg, .resp_msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}
            .user_msg {background-color:cornflowerblue;color:white;align-self:start}
            .resp_msg {background-color:lightgray;align-self:self-end}
        """,
        allow_flagging="never",
        theme="dark",
    )

    # launch the gradio interface and start the server
    iface.launch(
        enable_queue=True,
    )