| import os | |
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| from threading import Thread | |
| import torch | |
| tok = AutoTokenizer.from_pretrained("distilgpt2") | |
| model = AutoModelForCausalLM.from_pretrained("distilgpt2") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| n_gpu = 0 if torch.cuda.is_available()==False else torch.cuda.device_count() | |
| model.to(device) | |
| def generate(text = "", max_new_tokens = 128): | |
| streamer = TextIteratorStreamer(tok, timeout=10.) | |
| if len(text) == 0: | |
| text = " " | |
| inputs = tok([text], return_tensors="pt").to(device) | |
| generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=max_new_tokens, pad_token_id = model.config.eos_token_id, early_stopping=True, no_repeat_ngram_size=4) | |
| thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
| thread.start() | |
| generated_text = "" | |
| for new_text in streamer: | |
| yield generated_text + new_text | |
| generated_text += new_text | |
| if tok.eos_token in generated_text: | |
| generated_text = generated_text[: generated_text.find(tok.eos_token) if tok.eos_token else None] | |
| streamer.end() | |
| yield generated_text | |
| return | |
| return generated_text | |
| demo = gr.Interface( | |
| title="TextIteratorStreamer + Gradio demo", | |
| fn=generate, | |
| inputs=[gr.Textbox(lines=5, label="Input Text"), | |
| gr.Slider(default=128,minimum=5, maximum=256, step=1, label="Maximum number of new tokens")], | |
| outputs=gr.Textbox(label="Generated Text"), | |
| allow_flagging="never" | |
| ) | |
| demo.queue() | |
| demo.launch() |