File size: 1,639 Bytes
82ae8bf d149148 1ee42bd d149148 b18bc66 1ee42bd d149148 93c95b3 1ee42bd b18bc66 d149148 9e96195 d149148 0a3b879 9e96195 d149148 0a3b879 9e96195 d149148 8b3392b d4aa46c 8b3392b b18bc66 d149148 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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(value=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() |