Update app.py

app.py

@@ -1,29 +1,12 @@
-# 加载模型
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
-model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")
-# 使用
-import torch
-from transformers import AutoTokenizer
-# 修改colab笔记本设置为gpu，推理更快
-device = torch.device('cuda')
-model.to(device)
-def preprocess(text):
-  text = text.replace("\n", "\\n").replace("\t", "\\t")
-  return text
+import gradio as gr
 
-def postprocess(text):
-  return text.replace("\\n", "\n").replace("\\t", "\t")
+description = "Story generation with GPT-2"
+title = "Generate your own story"
+examples = [["Adventurer is approached by a mysterious stranger in the tavern for a new quest."]]
 
-def answer(text, sample=True, top_p=1, temperature=0.7):
-  '''sample：是否抽样。生成任务，可以设置为True;
-  top_p：0-1之间，生成的内容越多样'''
-  text = preprocess(text)
-  encoding = tokenizer(text=[text], truncation=True, padding=True, max_length=768, return_tensors="pt").to(device) 
-  if not sample:
-    out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, num_beams=1, length_penalty=0.6)
-  else:
-    out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, do_sample=True, top_p=top_p, temperature=temperature, no_repeat_ngram_size=3)
-  out_text = tokenizer.batch_decode(out["sequences"], skip_special_tokens=True)
-  return postprocess(out_text[0])
-print("end...")
+interface = gr.Interface.load("huggingface/pranavpsv/gpt2-genre-story-generator",
+            description=description,
+            examples=examples
+)
+
+interface.launch()