Create app.py

app.py

@@ -0,0 +1,45 @@
+!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
+    #print(new_text, end ="")
+    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()
+      #print("\n--\n")
+      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()