GLM-130B/server.py

import time
import torch
import torch.distributed as dist
import gradio as gr

from generation import BeamSearchStrategy, BaseStrategy
from initialize import initialize, initialize_model_and_tokenizer
from generate import add_generation_specific_args, fill_blanks


def generate_continually(func, raw_text):
    if not raw_text:
        return "Input should not be empty!"
    try:
        start_time = time.time()
        answer = func(raw_text)
        if torch.distributed.get_rank() == 0:
            print("\nTaken time {:.2f}\n".format(time.time() - start_time), flush=True)
        return answer
    except (ValueError, FileNotFoundError) as e:
        print(e)
        return "Error!"


strategy = None


def main(args):
    model, tokenizer = initialize_model_and_tokenizer(args)

    end_tokens = [tokenizer.get_command("eop"), tokenizer.get_command("eos")]

    def process(raw_text):
        global strategy

        if args.with_id:
            query_id, raw_text = raw_text.split("\t")

        answers, answers_with_style, blanks = fill_blanks(raw_text, model, tokenizer, strategy, args)

        if torch.distributed.get_rank() == 0:
            print(answers)

        return answers[0]

    def predict(
        text,
        seed=1234,
        out_seq_length=200,
        min_gen_length=20,
        sampling_strategy="BaseStrategy",
        num_beams=4,
        length_penalty=0.9,
        no_repeat_ngram_size=3,
        temperature=1,
        topk=1,
        topp=1,
    ):

        global strategy

        if torch.distributed.get_rank() == 0:
            print(
                "info",
                [
                    text,
                    seed,
                    out_seq_length,
                    min_gen_length,
                    sampling_strategy,
                    num_beams,
                    length_penalty,
                    no_repeat_ngram_size,
                    temperature,
                    topk,
                    topp,
                ],
            )
            dist.broadcast_object_list(
                [
                    text,
                    seed,
                    out_seq_length,
                    min_gen_length,
                    sampling_strategy,
                    num_beams,
                    length_penalty,
                    no_repeat_ngram_size,
                    temperature,
                    topk,
                    topp,
                ],
                src=0,
            )

        args.seed = seed
        args.out_seq_length = out_seq_length
        args.min_gen_length = min_gen_length
        args.sampling_strategy = sampling_strategy
        args.num_beams = num_beams
        args.length_penalty = length_penalty
        args.no_repeat_ngram_size = no_repeat_ngram_size
        args.temperature = temperature
        args.top_k = topk
        args.top_p = topp

        if args.sampling_strategy == "BaseStrategy":
            strategy = BaseStrategy(
                batch_size=1, temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, end_tokens=end_tokens
            )
        elif args.sampling_strategy == "BeamSearchStrategy":
            strategy = BeamSearchStrategy(
                batch_size=1,
                num_beams=args.num_beams,
                length_penalty=args.length_penalty,
                consider_end=True,
                end_tokens=end_tokens,
                no_repeat_ngram_size=args.no_repeat_ngram_size,
                min_gen_length=args.min_gen_length,
            )
        else:
            raise ValueError(f"unknown strategy {args.sampling_strategy}")

        return generate_continually(process, text)

    if torch.distributed.get_rank() == 0:
        en_fil = ["The Starry Night is an oil-on-canvas painting by [MASK] in June 1889."]
        en_gen = ["Eight planets in solar system are [gMASK]"]
        ch_fil = ["凯旋门位于意大利米兰市古城堡旁。1807年为纪念[MASK]而建，门高25米，顶上矗立两武士青铜古兵车铸像。"]
        ch_gen = ["三亚位于海南岛的最南端,是中国最南部的热带滨海旅游城市 [gMASK]"]
        en_to_ch = ["Pencil in Chinese is [MASK]."]
        ch_to_en = ['"我思故我在"的英文是"[MASK]"。']

        examples = [en_fil, en_gen, ch_fil, ch_gen, en_to_ch, ch_to_en]

        with gr.Blocks() as demo:
            gr.Markdown(
                """
                An Open Bilingual Pre-Trained Model. [Visit our github repo](https://github.com/THUDM/GLM-130B)
                GLM-130B uses two different mask tokens: `[MASK]` for short blank filling and `[gMASK]` for left-to-right long text generation. When the input does not contain any MASK token, `[gMASK]` will be automatically appended to the end of the text. We recommend that you use `[MASK]` to try text fill-in-the-blank to reduce wait time (ideally within seconds without queuing).
                """
            )

            with gr.Row():
                with gr.Column():
                    model_input = gr.Textbox(
                        lines=7, placeholder="Input something in English or Chinese", label="Input"
                    )
                    with gr.Row():
                        gen = gr.Button("Generate")
                        clr = gr.Button("Clear")

                outputs = gr.Textbox(lines=7, label="Output")

            gr.Markdown(
                """
                Generation Parameter
                """
            )
            with gr.Row():
                with gr.Column():
                    seed = gr.Slider(maximum=100000, value=1234, step=1, label="Seed")
                    out_seq_length = gr.Slider(
                        maximum=512, value=128, minimum=32, step=1, label="Output Sequence Length"
                    )
                with gr.Column():
                    min_gen_length = gr.Slider(maximum=64, value=0, step=1, label="Min Generate Length")
                    sampling_strategy = gr.Radio(
                        choices=["BeamSearchStrategy", "BaseStrategy"], value="BaseStrategy", label="Search Strategy"
                    )

            with gr.Row():
                with gr.Column():
                    # beam search
                    gr.Markdown(
                        """
                        BeamSearchStrategy
                        """
                    )
                    num_beams = gr.Slider(maximum=4, value=2, minimum=1, step=1, label="Number of Beams")
                    length_penalty = gr.Slider(maximum=1, value=1, minimum=0, label="Length Penalty")
                    no_repeat_ngram_size = gr.Slider(
                        maximum=5, value=3, minimum=1, step=1, label="No Repeat Ngram Size"
                    )
                with gr.Column():
                    # base search
                    gr.Markdown(
                        """
                        BaseStrategy
                        """
                    )
                    temperature = gr.Slider(maximum=1, value=1.0, minimum=0, label="Temperature")
                    topk = gr.Slider(maximum=40, value=0, minimum=0, step=1, label="Top K")
                    topp = gr.Slider(maximum=1, value=0.7, minimum=0, label="Top P")

            inputs = [
                model_input,
                seed,
                out_seq_length,
                min_gen_length,
                sampling_strategy,
                num_beams,
                length_penalty,
                no_repeat_ngram_size,
                temperature,
                topk,
                topp,
            ]
            gen.click(fn=predict, inputs=inputs, outputs=outputs)
            clr.click(fn=lambda value: gr.update(value=""), inputs=clr, outputs=model_input)

            gr_examples = gr.Examples(examples=examples, inputs=model_input)

        demo.launch(share=True)
    else:
        while True:
            info = [None, None, None, None, None, None, None, None, None, None, None]
            dist.broadcast_object_list(info, src=0)

            (
                text,
                seed,
                out_seq_length,
                min_gen_length,
                sampling_strategy,
                num_beams,
                length_penalty,
                no_repeat_ngram_size,
                temperature,
                topk,
                topp,
            ) = info

            predict(
                text,
                seed,
                out_seq_length,
                min_gen_length,
                sampling_strategy,
                num_beams,
                length_penalty,
                no_repeat_ngram_size,
                temperature,
                topk,
                topp,
            )


if __name__ == "__main__":
    args = initialize(extra_args_provider=add_generation_specific_args)

    with torch.no_grad():
        main(args)