|
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 |
|
|
|
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) |
|
|
|
early_stop_pattern = tok.eos_token |
|
print(f'Early stop pattern = \"{early_stop_pattern}\"') |
|
|
|
def generate(text = ""): |
|
streamer = TextIteratorStreamer(tok) |
|
if len(text) == 0: |
|
text = " " |
|
inputs = tok([text], return_tensors="pt") |
|
generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=128) |
|
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 early_stop_pattern in generated_text: |
|
generated_text = generated_text[: generated_text.find(early_stop_pattern) if early_stop_pattern else None] |
|
streamer.end() |
|
|
|
yield generated_text |
|
return |
|
|
|
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() |
