app.py

import streamlit as st
import requests
from enum import Enum


st.header("WeLM Demo 初体验")
st.text('Tips: ')
st.text("* WeLM 不是一个直接的对话机器人，而是一个补全用户输入信息的生成模型")
st.text("* 修改Prompts请先参考 https://welm.weixin.qq.com/docs/introduction/")

class Task(str, Enum):
    DIALOG_JOURNAL = "对话-采访"
    READING_COMPREHENSION = "阅读理解"
    TRANSLATE = "翻译"   
    REWRITE = "改写"
    COMPLETION = "续写"
    FREE = "自由任务"


task_value2type = {v.value: v.name for v in Task}

task_type = st.selectbox(
    "任务示例",
    [v.value for v in Task]
)
task_type = task_value2type[task_type]

task2prompt = {
    Task.READING_COMPREHENSION: """在这个关于“故意伤害罪”的案例：
    “经审理查明，被告人张××、杜×、杨2某均为辽宁省辽阳第一监狱五监区服刑人员。2015年11月3日13时许，被告人张××、杜×因无事便跟随去催要生产材料的被告人杨2某一同前往六监区，在六监区生产车间门外，被告人杨2某与六监区送料员于×因送料问题发生争执，被告人杨2某上前拽住被害人于×胳膊并用手击打被害人后脖颈两下，被告人张××、杜×见杨2某动手后，先后上前分别对被害人于×面部、头部及腹部进行殴打，后被赶到的干警制止。被害人于×被打造成面部受伤，鼻子流血，当日下午14时许，到监区内医院就诊，诊断为：鼻部中段向左侧畸形，11月5日经监狱医院X光诊断为鼻骨骨折。2015年11月18日，经辽阳襄平法医司法鉴定所法医鉴定：被害人于×身体损伤程度为轻伤二级。被告人张××、杜×、杨2某共同赔偿被害人于×人民币7000元，被害人于×对被告人的行为表示谅解。”
问题: “被害人于×11月5日经监狱医院X光诊断后的诊断结果为？”
答案：""",
    Task.TRANSLATE: """“I had a good time”的中文翻译是：我玩得很开心。
“The pandemic has triggered more demand for online shopping”的中文翻译是：疫情引发了更多的网购需求
“I am a programmer in Tencent”的中文翻译是：""",
    Task.DIALOG_JOURNAL: """Elon Musk 是 SpaceX 和特斯拉汽⻋公司的首席执行官。
我：今天我们有 Elon Musk。欢迎加入我们。
Elon Musk：谢谢你邀请我。
我：你认为 OpenAI 会如何发展？
Elon Musk：我认为进展非常顺利。我认为我们在 OpenAI 有一个非常有才华的团队。
我：你现在是怎么进行时间分配？
Elon Musk：我每周都会尝试将部分时间花在 OpenAI 上，但除此之外，其他时间是 SpaceX 和特斯拉。
我：你如何看待自动驾驶技术？
Elon Musk：""",
    Task.REWRITE: """有这样一段文本，{医生微笑着递给小明棒棒糖，同时让小明服下了药。}
改写这段话让它变得更加惊悚。{医生眼露凶光让小明服药，小明感到非常害怕}。
有这样一段文本，{雨下得很大}
改写这段话让它变得更加具体。{一霎时，雨点连成了线，大雨就像天塌了似的铺天盖地从空中倾泻下来。}。
有这样一段文本，{王老师离开了电影院，外面已经天黑了}
改写这段话让它包含更多电影信息。{这部电影比小王预想的时间要长，虽然口碑很好，但离开电影院时，小王还是有些失望。}
有这样一段文本，{男人站在超市外面打电话}
改写这段话来描述小丑。{男人站在马戏团外一边拿着气球一边打电话}
有这样一段文本，{风铃声响起}
改写这段话写的更加丰富。{我对这个风铃的感情是由它的铃声引起的。每当风吹来时，风铃发出非常动听的声音，听起来是那么乐观、豁达，像一个小女孩格格的笑声。}
""",
    Task.COMPLETION: """“八月十八潮，壮观天下无。”这是北宋大诗人苏东坡咏赞钱塘秋潮的千古名句。千百年来，钱塘江以其奇特卓绝的江潮，不知倾倒了多少游人看客。
每年的农历八月十八前后，是观潮的最佳时节。这期间，秋阳朗照，金风宜人，钱塘江口的海塘上，游客群集，兴致盎然，争睹奇景。
""",
    Task.FREE: ""
}

prompt = task2prompt[Task[task_type]]
if task_type == 'REWRITE':
    pre_input = st.text_area('输入样例', value=prompt, height=400)
    rewrite_text = "有这样一段文本，{我想家了}\n改写这段话包含更多悲伤的感情。{"
    rewrite_text = st.text_area('你的输入', value=rewrite_text, height=100)
    all_input = prompt + rewrite_text  
else:
    all_input = st.text_area('模型输入', value=prompt, height=400)
all_input = all_input.rstrip('\\n')


with st.expander("配置"):
    stop_tokens = ""
    def cut_message(answer: str):
        end = []
        for etk in stop_tokens:
            offset = answer.find(etk)
            if offset > 0:
                end.append(offset)
        if len(end) > 0:
            answer = answer[:min(end)]
        return answer.rstrip()

    if task_type == 'READING_COMPREHENSION':
        default_top_p, default_temperature, default_n, default_tokens = 0.0, 0.0, 1, 15
    elif task_type == 'TRANSLATE':
        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
    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)
    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)
    max_tokens = st.slider('max tokens', 4, 512, default_tokens)
    
    if st.checkbox("使用换行符作为截断", value=True):
        stop_tokens = "\n"

def completion():
    try:
        resp = requests.post("https://welm.weixin.qq.com/v1/completions", json={
            'prompt': all_input,
            'max_tokens': max_tokens,
            'temperature': temperature,
            'top_p': top_p,
            'top_k': top_k,
            'n': n,
            'model': model,
            "stop": stop_tokens,
        }, headers={"Authorization": f"Bearer {st.secrets['token']}"})
        answer = resp.json()
        for idx, choice in enumerate(answer['choices']):
            if choice.get("finish_reason", None) != "finished":
                st.error(f'生成结果#{idx}出错: {choice["finish_reason"]}')
            elif choice.get("text", None) is None:
                st.error(f'生成结果#{idx}出错: 无生成结果')
            elif len(choice.get("text", "")) == 0:
                st.info(f'生成结果#{idx}: 生成结果为空')
            else:    
                st.success(f'生成结果#{idx}: {cut_message(choice["text"])}')
        
        if task_type == 'COMPLETION':
            st.text('Tips: 可多次生成后复制你认为的最好结果拼接于原文后，让WeLM继续生成。')

    except Exception as e:
        st.error(f"生成结果出错：{str(e)}")
    
   

if st.button('立即生成'):
    completion()


footer="""<style>

a:link , a:visited{
background-color: transparent;
text-decoration: none;
color: rgb(7 6 17);
}

a:hover,  a:active {
background-color: transparent;
text-decoration: underline;
}

.footer {
position: relative;
left: 0;
bottom: 0;
width: 100%;
background-color: white;
color: rgb(18 19 23 / 80%);
text-align: left;
}
</style>
<div class="footer" style="font-size:12px">
<br>
<br>
<br>
<br> 
<br> 
<p style="font-size:12px">声明（受<a href="https://beta.openai.com/docs/usage-guidelines/content-policy" target="_blank"> OpenAI GPT3's Content Policy </a>启发）</p>

<p style="font-size:12px">我们禁止用户在知情的情况下产生或利用他人在知情的情况下产生有害内容，包括且不限于仇恨言论、骚扰信息、自我伤害内容、性内容、政治内容、垃圾邮件、诈骗信息等。<br> 
特别提示：本网页上产生的所有内容不代表本平台的观点和意见。<br> 
欲了解更多信息，请阅读我们的<a href="https://welm.weixin.qq.com/docs/license/" target="_blank">《服务使用协议》和《免责声明》</a>。</p>
</div>
"""
st.markdown(footer,unsafe_allow_html=True)