Update app.py

app.py

@@ -8,11 +8,12 @@ st.text('Tips: ')
 st.text("* WeLM不是一个直接的对话机器人，而是一个补全用户输入信息的生成模型")
 st.text("* 修改Prompt可以更多参考 https://welm.weixin.qq.com/docs/introduction/")
 st.text("* 你的输入可能会被我们拼接在预设的prompt尾部后再发送给API")
-st.text("* 在每个任务的下方我们展示了该任务请求API时完整的参数（包含完整的prompt）")
+# st.text("* 在每个任务的下方我们展示了该任务请求API时完整的参数（包含完整的prompt）")
 
 
 
 class Task(str, Enum):
+    FREE = "自由任务"
     DIALOG_JOURNAL = "对话（Elon musk）"
     QA = "问答"
     COPY= "文案生成"
@@ -20,7 +21,7 @@ class Task(str, Enum):
     READING_COMPREHENSION = "阅读理解"
     TRANSLATE = "翻译"   
     COMPLETION = "文章续写"    
-    FREE = "自由任务"
+
 
 
 task_value2type = {v.value: v.name for v in Task}
@@ -32,6 +33,17 @@ task_type = st.selectbox(
 task_type = task_value2type[task_type]
 
 task2prompt_pre = {
+    Task.FREE: """
+
+Human: 你是谁<eot>
+
+Assistant: 我是 Assistant，一个由 WeChat 训练的 AI 语言模型。
+我对人类活动有全面的了解，可以为你解答各种领域的问题。
+我会在最短的篇幅内给你信息量最丰富的回答。
+我的回答是理性、准确的。
+在我认为合适的时候，我可能会用英文回复你。<eot>
+
+Human: """,
     Task.READING_COMPREHENSION: """请阅读文章后根据文章内容给出问题的答案。
 文章：中国空间技术研究院（China Academy of Space Technology，简称CAST）隶属于中国航天科技集团公司，是中国空间技术的主要研究中心和航天器研制、生产基地，成立于1968年2月20日。下设10个研究所和一个工厂。现任院长为杨保华，院党委书记为李开民。1970年4月24日，中国空间技术研究院成功研制并发射了中国首颗人造地球卫星东方红一号。2003年10月，神舟五号载人飞船载人飞行取得成功。2005年，神舟六号载人飞船实现多人多天的太空飞行。截至2005年，中国空间技术研究院研制并成功发射了68颗不同类型的人造卫星、4艘无人试验飞船和2艘载人飞船，涵盖通信广播卫星、返回式遥感>卫星、地球资源卫星、气象卫星、科学探测与技术试验卫星、导航定位卫星和载人航天器等领域。
 问题：中国空间技术研究院在哪年成立？
@@ -89,7 +101,6 @@ Elon Musk: 51岁，但我说我有40岁人的精力。在健康方面最重要
 """,
     Task.COMPLETION: """
 """,
-    Task.FREE: ""
 }
 
 task2prompt_end = {
@@ -114,13 +125,15 @@ Elon Musk：""",
 prompt_fix = task2prompt_pre[Task[task_type]]
 prompt_user = task2prompt_end[Task[task_type]]
 
-user_input = st.text_area('你的输入（最终完整输入请见下方 API 请求内容）', value=prompt_user, height=180)
+user_input = st.text_area('你的输入', value=prompt_user, height=150)
 all_input = prompt_fix + user_input
+if Task[task_type] == Task.FREE:
+    all_input = prompt_fix + user_input + "<eot>\n\nAssistant: "
 all_input = all_input.rstrip('\\n')
 
 
 with st.expander("配置"):
-    stop_tokens = ""
+    stop_tokens = ["<eot>", "\n\nHuman"]
     def cut_message(answer: str):
         end = []
         for etk in stop_tokens:
@@ -137,14 +150,22 @@ with st.expander("配置"):
         default_top_p, default_temperature, default_n, default_tokens = 0.0, 0.0, 1, 60       
     elif task_type == 'COMPLETION':
         default_top_p, default_temperature, default_n, default_tokens = 0.95, 0.85, 1, 150
+    elif task_type == 'FREE':
+        default_top_p, default_temperature, default_n, default_tokens = 0.95, 0.85, 1, 300
     else:
         default_top_p, default_temperature, default_n, default_tokens = 0.95, 0.85, 3, 64
     
-    model = st.selectbox("model", ["medium", "large", "xl"], index=2)
+    if task_type == 'FREE':
+        model_list = ["xl-answer"]
+    else:
+        model_list = ["xl", "large", "medium"]
+    model = st.selectbox("model", model_list, index=0)
     top_p = st.slider('top p', 0.0, 1.0, default_top_p)
     top_k = st.slider('top k', 0, 100, 0)
     temperature = st.slider('temperature', 0.0, 1.0, default_temperature)
     n = st.slider('n', 1, 5, default_n)
+    if task_type == 'FREE':
+        n = 1
     max_tokens = st.slider('max tokens', 4, 512, default_tokens)
     
     if st.checkbox("使用换行符作为截断", value=False):
@@ -183,19 +204,19 @@ def completion():
         st.error(f"生成结果出错：{str(e)}")
     
    
-    code_str = """
-        post_json = {{
-        'prompt': '{all_input}',
-        'model': '{model}',
-        'max_tokens': {max_tokens},
-        'temperature': {temperature},
-        'top_p': {top_p},
-        'top_k': {top_k},
-        'n': {n},
-        "stop": '{stop_tokens}',
-    }}  
-    """.format(all_input=all_input,model=model,max_tokens=max_tokens,temperature=temperature, top_p=top_p,top_k=top_k,n=n,stop_tokens=stop_tokens)
-    st.code(code_str)
+    # code_str = """
+    #     post_json = {{
+    #     'prompt': '{all_input}',
+    #     'model': '{model}',
+    #     'max_tokens': {max_tokens},
+    #     'temperature': {temperature},
+    #     'top_p': {top_p},
+    #     'top_k': {top_k},
+    #     'n': {n},
+    #     "stop": '{stop_tokens}',
+    # }}  
+    # """.format(all_input=all_input,model=model,max_tokens=max_tokens,temperature=temperature, top_p=top_p,top_k=top_k,n=n,stop_tokens=stop_tokens)
+    # st.code(code_str)
 
 if st.button('立即生成'):
     completion()