Update app.py

app.py

@@ -4,6 +4,7 @@ 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")
@@ -16,10 +17,10 @@ early_stop_pattern = tok.eos_token
 print(f'Early stop pattern = \"{early_stop_pattern}\"')
 
 def generate(text = ""):  
-  streamer = TextIteratorStreamer(tok)
+  streamer = TextIteratorStreamer(tok, timeout=10.)
   if len(text) == 0:
     text = " "
-  inputs = tok([text], return_tensors="pt")  
+  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)
   thread = Thread(target=model.generate, kwargs=generation_kwargs)
   thread.start()