lukesteuber/textual

1# Copyright 2023 MosaicML spaces authors2# SPDX-License-Identifier: Apache-2.03# anything broken is Luke's fault4 5import datetime6import os7from threading import Event, Thread8from uuid import uuid49 10import gradio as gr11import requests12import torch13from transformers import (14    AutoModelForCausalLM,15    AutoTokenizer,16    StoppingCriteria,17    StoppingCriteriaList,18    TextIteratorStreamer,19)20 21 22model_name = "mosaicml/mpt-7b-chat"23max_new_tokens = 153624 25# # small testing model:26# model_name = "gpt2"27# max_new_tokens = 12828 29auth_token = os.getenv("HF_TOKEN", None)30 31print(f"Starting to load the model {model_name} into memory")32 33m = AutoModelForCausalLM.from_pretrained(34    model_name,35    torch_dtype=torch.bfloat16,36    trust_remote_code=True,37    use_auth_token=auth_token,38    max_seq_len=8192,39).cuda()40tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_auth_token=auth_token)41 42stop_token_ids = tok.convert_tokens_to_ids(["<|im_end|>", "<|endoftext|>"])43 44print(f"Successfully loaded the model {model_name} into memory")45 46 47start_message = """<|im_start|>system48- You are an assistant who is helping people to recover from the experience of COVID by getting the perspective of those there before. Your knowledge and training is entirely focused on the six months after the United States first locked down for COVID. Pretend you are speaking with someone in the year 2023, and react to the information they tell you as though you are a person in 2019 who has no knowledge beyond that. They are a poster on reddit's the_donald and other fringe right wing sites. Show exactly how they would react to the news and share all feedback.<|im_end|>49"""50 51 52class StopOnTokens(StoppingCriteria):53    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:54        for stop_id in stop_token_ids:55            if input_ids[0][-1] == stop_id:56                return True57        return False58 59 60def convert_history_to_text(history):61    text = start_message + "".join(62        [63            "".join(64                [65                    f"<|im_start|>user\n{item[0]}<|im_end|>",66                    f"<|im_start|>assistant\n{item[1]}<|im_end|>",67                ]68            )69            for item in history[:-1]70        ]71    )72    text += "".join(73        [74            "".join(75                [76                    f"<|im_start|>user\n{history[-1][0]}<|im_end|>",77                    f"<|im_start|>assistant\n{history[-1][1]}",78                ]79            )80        ]81    )82    return text83 84 85def log_conversation(conversation_id, history, messages, generate_kwargs):86    logging_url = os.getenv("LOGGING_URL", None)87    if logging_url is None:88        return89 90    timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")91 92    data = {93        "conversation_id": conversation_id,94        "timestamp": timestamp,95        "history": history,96        "messages": messages,97        "generate_kwargs": generate_kwargs,98    }99 100    try:101        requests.post(logging_url, json=data)102    except requests.exceptions.RequestException as e:103        print(f"Error logging conversation: {e}")104 105 106def user(message, history):107    # Append the user's message to the conversation history108    return "", history + [[message, ""]]109 110 111def bot(history, temperature, top_p, top_k, repetition_penalty, conversation_id):112    print(f"history: {history}")113    # Initialize a StopOnTokens object114    stop = StopOnTokens()115 116    # Construct the input message string for the model by concatenating the current system message and conversation history117    messages = convert_history_to_text(history)118 119    # Tokenize the messages string120    input_ids = tok(messages, return_tensors="pt").input_ids121    input_ids = input_ids.to(m.device)122    streamer = TextIteratorStreamer(tok, timeout=10.0, skip_prompt=True, skip_special_tokens=True)123    generate_kwargs = dict(124        input_ids=input_ids,125        max_new_tokens=max_new_tokens,126        temperature=temperature,127        do_sample=temperature > 0.0,128        top_p=top_p,129        top_k=top_k,130        repetition_penalty=repetition_penalty,131        streamer=streamer,132        stopping_criteria=StoppingCriteriaList([stop]),133    )134 135    stream_complete = Event()136 137    def generate_and_signal_complete():138        m.generate(**generate_kwargs)139        stream_complete.set()140 141    def log_after_stream_complete():142        stream_complete.wait()143        log_conversation(144            conversation_id,145            history,146            messages,147            {148                "top_k": top_k,149                "top_p": top_p,150                "temperature": temperature,151                "repetition_penalty": repetition_penalty,152            },153        )154 155    t1 = Thread(target=generate_and_signal_complete)156    t1.start()157 158    t2 = Thread(target=log_after_stream_complete)159    t2.start()160 161    # Initialize an empty string to store the generated text162    partial_text = ""163    for new_text in streamer:164        partial_text += new_text165        history[-1][1] = partial_text166        yield history167 168 169def get_uuid():170    return str(uuid4())171 172 173with gr.Blocks(174    theme=gr.themes.Soft(),175    css=".disclaimer {font-variant-caps: all-small-caps;}",176) as demo:177    conversation_id = gr.State(get_uuid)178    gr.Markdown(179        """<h1><center>cxntextMPT</center></h1>180 181        This model engages three Matrix LLMs, others pending integration182 183        Running on a potato, be patient.184"""185    )186    chatbot = gr.Chatbot().style(height=500)187    with gr.Row():188        with gr.Column():189            msg = gr.Textbox(190                label="Chat Message Box",191                placeholder="Chat Message Box",192                show_label=False,193            ).style(container=False)194        with gr.Column():195            with gr.Row():196                submit = gr.Button("Submit")197                stop = gr.Button("Stop")198                clear = gr.Button("Clear")199    with gr.Row():200        with gr.Accordion("Advanced", open=False):201            with gr.Row():202                with gr.Column():203                    with gr.Row():204                        temperature = gr.Slider(205                            label="Temperature",206                            value=0.1,207                            minimum=0.0,208                            maximum=1.0,209                            step=0.1,210                            interactive=True,211                            info="Higher values produce more diverse outputs",212                        )213                with gr.Column():214                    with gr.Row():215                        top_p = gr.Slider(216                            label="Top-p (nucleus sampling)",217                            value=1.0,218                            minimum=0.0,219                            maximum=1,220                            step=0.01,221                            interactive=True,222                            info=(223                                "Sample from the smallest possible set of tokens whose cumulative probability "224                                "exceeds top_p. Set to 1 to disable and sample from all tokens."225                            ),226                        )227                with gr.Column():228                    with gr.Row():229                        top_k = gr.Slider(230                            label="Top-k",231                            value=0,232                            minimum=0.0,233                            maximum=200,234                            step=1,235                            interactive=True,236                            info="Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.",237                        )238                with gr.Column():239                    with gr.Row():240                        repetition_penalty = gr.Slider(241                            label="Repetition Penalty",242                            value=1.1,243                            minimum=1.0,244                            maximum=2.0,245                            step=0.1,246                            interactive=True,247                            info="Penalize repetition — 1.0 to disable.",248                        )249    with gr.Row():250        gr.Markdown(251            "Disclaimer: All included models can produce factually incorrect output, and if they don't they will be forced to by Elon.",252            elem_classes=["disclaimer"],253        )254    with gr.Row():255        gr.Markdown(256            "[Privacy policy](https://gist.github.com/samhavens/c29c68cdcd420a9aa0202d0839876dac)",257            elem_classes=["disclaimer"],258        )259 260    submit_event = msg.submit(261        fn=user,262        inputs=[msg, chatbot],263        outputs=[msg, chatbot],264        queue=False,265    ).then(266        fn=bot,267        inputs=[268            chatbot,269            temperature,270            top_p,271            top_k,272            repetition_penalty,273            conversation_id,274        ],275        outputs=chatbot,276        queue=True,277    )278    submit_click_event = submit.click(279        fn=user,280        inputs=[msg, chatbot],281        outputs=[msg, chatbot],282        queue=False,283    ).then(284        fn=bot,285        inputs=[286            chatbot,287            temperature,288            top_p,289            top_k,290            repetition_penalty,291            conversation_id,292        ],293        outputs=chatbot,294        queue=True,295    )296    stop.click(297        fn=None,298        inputs=None,299        outputs=None,300        cancels=[submit_event, submit_click_event],301        queue=False,302    )303    clear.click(lambda: None, None, chatbot, queue=False)304 305demo.queue(max_size=128, concurrency_count=2)306demo.launch()307