Spaces:
Runtime error
Runtime error
| import json | |
| from threading import Thread | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer | |
| ##### MODEL AND TOKENIZER SETUP ###### | |
| model_id = "mylesmharrison/gpt2-moviedialog" | |
| model = AutoModelForCausalLM.from_pretrained(model_id, from_tf=True) | |
| tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| ##### TEXT GENERATION (INFERENCE) ###### | |
| def run_generation(user_text, top_p, temperature, top_k, max_new_tokens): | |
| # Get the model and tokenizer, and tokenize the user text. | |
| model_inputs = tokenizer([user_text], return_tensors="pt") | |
| # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer | |
| # in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread. | |
| streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True) | |
| generate_kwargs = dict( | |
| model_inputs, | |
| streamer=streamer, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=True, | |
| top_p=top_p, | |
| temperature=float(temperature), | |
| top_k=top_k | |
| ) | |
| t = Thread(target=model.generate, kwargs=generate_kwargs) | |
| t.start() | |
| # Pull the generated text from the streamer, and update the model output. | |
| model_output = "" | |
| for new_text in streamer: | |
| model_output += new_text | |
| yield model_output | |
| return model_output | |
| ##### HELPER FUNCTION AND TEXT DATA ###### | |
| def reset_textbox(): | |
| return gr.update(value='') | |
| # Read in the text data for examples | |
| with open("gradio_data.json", "r") as f: | |
| gr_data = json.loads(f.read()) | |
| # Read in the markdown data for the header | |
| with open("README.md", "r") as f: | |
| desc = ''.join(f.readlines()[11:]) | |
| theme = gr.themes.Base() | |
| ##### UI RENDER ###### | |
| with gr.Blocks(theme=theme) as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| ### DESCRIPTION | |
| gr.Markdown(desc) | |
| with gr.Row(): | |
| with gr.Column(): | |
| # INPUT FIELD | |
| user_text = gr.Textbox( | |
| placeholder="JOHN: I love you.", | |
| label="User input", | |
| scale=4, | |
| lines=10 | |
| ) | |
| # SUBMIT BUTTON | |
| with gr.Row(): | |
| with gr.Column(): | |
| button_submit = gr.Button(value="Submit") | |
| with gr.Column(): | |
| stop = gr.Button(value="Stop") | |
| ### TEXT EXAMPLES | |
| gr.Examples(gr_data['examples'],[user_text]) | |
| with gr.Column(): | |
| with gr.Row(scale=4): | |
| model_output = gr.Textbox(label="Model output", lines=10, interactive=False) | |
| with gr.Row(scale=1): | |
| ### INFERENCE CONTROLS | |
| max_new_tokens = gr.Slider( | |
| minimum=1, maximum=500, value=150, step=10, interactive=True, label="Max Tokens", | |
| ) | |
| top_p = gr.Slider( | |
| minimum=0.05, maximum=1.0, value=0.95, step=0.05, interactive=True, label="Top-p", | |
| ) | |
| top_k = gr.Slider( | |
| minimum=1, maximum=50, value=20, step=1, interactive=True, label="Top-k", | |
| ) | |
| temperature = gr.Slider( | |
| minimum=0.1, maximum=5.0, value=1.1, step=0.1, interactive=True, label="Temperature", | |
| ) | |
| # RUN | |
| text_event = user_text.submit(run_generation, [user_text, top_p, temperature, top_k, max_new_tokens], model_output) | |
| button_event = button_submit.click(run_generation, [user_text, top_p, temperature, top_k, max_new_tokens], model_output) | |
| stop.click(fn=None, inputs=None, outputs=None, cancels=[text_event, button_event]) | |
| demo.queue(max_size=32).launch(enable_queue=True) |