Initial app.py

app.py

@@ -21,5 +21,36 @@
 
 
 import gradio as gr
+from transformers import BertTokenizer, BartForConditionalGeneration
 
-gr.Interface.load("models/HIT-TMG/dialogue-bart-large-chinese").launch()
+
+title = "标题"
+description = "描述"
+
+
+tokenizer = BertTokenizer.from_pretrained("HIT-TMG/dialogue-bart-large-chinese")
+model = BartForConditionalGeneration.from_pretrained("HIT-TMG/dialogue-bart-large-chinese")
+
+
+def chat(history):
+    history_prefix = "对话历史："
+    history = history_prefix + history
+
+    input_ids = tokenizer(history, return_tensors='pt').input_ids
+    output_ids = model.generate(input_ids)[0]
+
+    return tokenizer.decode(output_ids, skip_special_tokens=True)
+
+
+chatbot = gr.Chatbot().style(color_map=("green", "pink"))
+demo = gr.Interface(
+    chat,
+    inputs=gr.Textbox(lines=2, placeholder="输入你的对话历史"),
+    title=title,
+    description=description,
+    outputs =["text"]
+)
+
+
+if __name__ == "__main__":
+    demo.launch()