| import os | |
| os.system("pip install git+https://github.com/huggingface/transformers") | |
| 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 = ""): | |
| 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=128, 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.inputs.Textbox(lines=5, label="Input Text"), | |
| outputs=gr.outputs.Textbox(label="Generated Text"), | |
| ) | |
| demo.queue() | |
| demo.launch() |