Update app.py

app.py

@@ -2,14 +2,19 @@ import gradio as gr
 from transformers import pipeline
 import random
 import re
+import torch
+import spaces
 
 title = "WoW Quest Text Generator"
 description = "Tap on the \"Submit\" button to generate a random quest text."
 article = "<p>Fine tuned <a href=\"https://huggingface.co/EleutherAI/gpt-neo-125M\">EleutherAI/gpt-neo-125M</a> upon a formatted <a href=\"https://github.com/TrinityCore/TrinityCore\"> TrinityCore – TDB_full_world_927.22082_2022_08_21 Dataset</a></p><p>This generator is fan made and is not affiliated in any way with Blizzard and/or any other company</p>"
 
+CUDA_AVAILABLE = torch.cuda.is_available()
+device = torch.device("cuda" if CUDA_AVAILABLE else "cpu")
+
 model_id = "./model"
-text_generator = pipeline("text-generation", model=model_id, tokenizer=model_id)
-max_length = 192
+text_generator = pipeline("text-generation", model=model_id, tokenizer=model_id, device=device)
+max_length = 256
 top_k = 40
 top_p = 0.92
 temperature = 1.0
@@ -43,6 +48,7 @@ def parseSpecialCharacters(text, wow_class_item, wow_race_item, wow_silly_name_i
     parsedText = text.replace("$a", "\n").replace("$B", "\n").replace("$b", "\n").replace("$c", wow_class_item).replace("$C", wow_class_item).replace("$r", wow_race_item).replace("$R", wow_race_item).replace("$n", wow_silly_name_item).replace("$N", wow_silly_name_item)
     return parseGenderTokens(parsedText)
 
+@spaces.GPU
 def text_generation(input_text = None):
     if input_text == None or len(input_text) == 0:
         input_text = "<|startoftext|>"
@@ -65,14 +71,17 @@ def text_generation(input_text = None):
     parsed_text = parsed_text.replace("\\n", "\n")
     return parsed_text
 
-gr.Interface(
+demo = gr.Interface(
     text_generation,    
-    [gr.inputs.Textbox(lines=1, label="Enter strating text or leave blank")],
-    outputs=[gr.outputs.Textbox(type="text", label="Generated quest text")],
+    inputs=gr.Textbox(Textbox(lines=1, label="Enter strating text or leave blank"),
+    outputs=gr.Textbox(type="text", label="Generated quest text")],
     title=title,
     description=description,
     article=article,
     examples=suggested_text_list,
     theme="default",
     allow_flagging=False,
-).launch()
+)
+
+demo.queue()
+demo.launch()