Update app.py

app.py

@@ -33,6 +33,7 @@ from transformers import MegatronBertForMaskedLM
 import argparse
 import copy
 import streamlit as st
+import time
 # os.environ["CUDA_VISIBLE_DEVICES"] = '6'
 
 
@@ -612,12 +613,12 @@ def comp_acc(pred_data, test_data):
 
 
 @st.experimental_memo()
-def load_model():
+def load_model(model_parh):
     total_parser = argparse.ArgumentParser("TASK NAME")
     total_parser = UniMCPipelines.pipelines_args(total_parser)
     args = total_parser.parse_args()
 
-    args.pretrained_model_path = 'IDEA-CCNL/Erlangshen-UniMC-RoBERTa-110M-Chinese'
+    args.pretrained_model_path = model_path
     args.max_length = 512
     args.batchsize = 8
     args.default_root_dir = './'
@@ -628,14 +629,52 @@ def load_model():
 
 def main():
 
-    model = load_model()
+    text_dict={
+        '文本分类':"微软披露拓扑量子计算机计划！",
+        '情感分析':"刚买iphone13 pro 还不到一个月，天天死机最差的一次购物体验",
+        '语义匹配':"今天心情不好,我很不开心",
+        '自然语言推理':"小明正在上高中[unused1]小明是一个初中生",
+        '多项式阅读理解':"这个男的是什么意思?[unused1][SEP]女：您看这件衣服挺不错的，质量好，价钱也不贵。\n男：再看看吧。",
+    }
+
+    question_dict={
+        '文本分类':"故事；文化；娱乐；体育；财经；房产；汽车；教育；科技",
+        '情感分析':"好评；差评",
+        '语义匹配':"可以理解为；不能理解为",
+        '自然语言推理':"可以推断出；不能推断出；很难推断出",
+        '多项式阅读理解':"不想要这件；衣服挺好的；衣服质量不好",
+    }
+
+    choice_dict={
+        '文本分类':"故事；文化；娱乐；体育；财经；房产；汽车；教育；科技",
+        '情感分析':"好评；差评",
+        '语义匹配':"可以理解为；不能理解为",
+        '自然语言推理':"可以推断出；不能推断出；很难推断出",
+        '多项式阅读理解':"不想要这件；衣服挺好的；衣服质量不好",
+    }
+    
+    
 
     st.subheader("UniMC Zero-shot 体验")
+    
+    st.sidebar.header("参数配置")
+    sbform = st.sidebar.form("固定参数设置")
+    language = sbform.selectbox('选择语言', ['中文', 'English'])
+    sbform.form_submit_button("配置")
+    
+    if language == '中文':
+        model = load_model('IDEA-CCNL/Erlangshen-UniMC-RoBERTa-110M-Chinese')
+    else:
+        model = load_model('IDEA-CCNL/Erlangshen-UniMC-RoBERTa-110M-Chinese')
+    
+    model_type = st.selectbox('选择任务类型',['文本分类','情感分析','语义匹配','自然语言推理','多项式阅读理解'])
+    
     st.info("请输入以下信息...")
+    
 
-    sentences = st.text_area("请输入句子：", """彭于晏不着急，胡歌也不着急，他俩都不着急，那我也不着急""")
-    question = st.text_input("请输入问题（不输入问题也可以）：", "请问下面的新闻属于哪个类别？")
-    choice = st.text_input("输入标签（以中文；分割）:", "娱乐；军事；体育；财经")
+    sentences = st.text_area("请输入句子：", text_dict[model_type])
+    question = st.text_input("请输入问题（不输入问题也可以）：", "")
+    choice = st.text_input("输入标签（以中文；分割）:", choice_dict[model_type])
     choice = choice.split('；')
 
     data = [{"texta": sentences,
@@ -646,8 +685,9 @@ def main():
             "id": 0}]
 
     if st.button("点击一下，开始预测！"):
+        start=time.time()
         result = model.predict(data, cuda=False)
-        st.success("预测成功！")
+        st.success(f"Prediction is successful, consumes {str(time.time()-start)} seconds")
         st.json(result[0])
     else:
         st.info(