import streamlit as st
import requests
from enum import Enum
st.header("WeLM Demo 初体验")
st.text('Tips: ')
st.text("* WeLM不是一个直接的对话机器人，而是一个补全用户输入信息的生成模型")
st.text("* 修改Prompt可以更多参考 https://welm.weixin.qq.com/docs/introduction/")
st.text("* 你的输入可能会被我们拼接在预设的prompt尾部后再发送给API")
st.text("* 在每个任务的下方我们展示了该任务请求API时完整的参数（包含完整的prompt）")
class Task(str, Enum):
DIALOG_JOURNAL = "对话（Elon musk）"
QA = "问答"
COPY= "文案生成"
REWRITE = "文本改写"
READING_COMPREHENSION = "阅读理解"
TRANSLATE = "翻译"
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_pre = {
Task.READING_COMPREHENSION: """请阅读文章后根据文章内容给出问题的答案。
文章：中国空间技术研究院（China Academy of Space Technology，简称CAST）隶属于中国航天科技集团公司，是中国空间技术的主要研究中心和航天器研制、生产基地，成立于1968年2月20日。下设10个研究所和一个工厂。现任院长为杨保华，院党委书记为李开民。1970年4月24日，中国空间技术研究院成功研制并发射了中国首颗人造地球卫星东方红一号。2003年10月，神舟五号载人飞船载人飞行取得成功。2005年，神舟六号载人飞船实现多人多天的太空飞行。截至2005年，中国空间技术研究院研制并成功发射了68颗不同类型的人造卫星、4艘无人试验飞船和2艘载人飞船，涵盖通信广播卫星、返回式遥感>卫星、地球资源卫星、气象卫星、科学探测与技术试验卫星、导航定位卫星和载人航天器等领域。
问题：中国空间技术研究院在哪年成立？
答案：1968年
""",
Task.QA: """请根据所学知识回答问题
问题：百年孤独的作者是？
回答：作者是哥伦比亚作家加西亚·马尔克斯，这本书是其代表作，也是拉丁美洲魔幻现实主义文学的代表作，被誉为“再现拉丁美洲历史社会图景的鸿篇巨著”。
问题：世界第二高峰是？
回答：乔戈里峰。海拔8611，海拔仅次于珠穆朗玛峰。“乔戈里”，通常被认为是塔吉克语，意思是“高大雄伟”。乔戈里山峰主要有6条山脊，西北—东南山脊为喀喇昆山脉主脊线，同时也是中国、巴基斯坦的国境线。
""",
Task.COPY: """请根据商品描述生成商品的文案
商品描述:芍药香氛的沐浴乳
文案:无比高大上的香味，淡粉色玫瑰清新诱人！沐浴后都充满着浪漫和幸福感，这样的情调会让你变得更加温柔。
商品描述:清爽去痘的洗面奶
文案:蕴含海藻精华，配合多种草本植物。能清洁毛孔污垢，保湿滋润肌肤，让细胞更健康、平衡水油分泌，消除痘痘同时预防痘痘产生，预防黑头粉刺及暗疮形成。""",
Task.TRANSLATE: """“I had a good time”的中文翻译是：我玩得很开心。
“The pandemic has triggered more demand for online shopping”的中文翻译是：疫情引发了更多的网购需求
""",
Task.DIALOG_JOURNAL: """以下是与Elon Musk的一系列对话。很多人都知道，Elon Musk是特斯拉、SpaceX、Neuralink和Boring公司的领导人。可能鲜为人知的是，他是一位世界级的工程师和设计师，不断强调第一原理的思考，并承担了许多在他之前的人认为不可能的大工程问题。Elon Musk 在2022年4月达成了一项有争议的交易，以440亿美元收购了twitter。经过无数次商业和法律上的对抗，收购于10月27日正式完成。收购后，马斯克立即解雇了公司的大部分员工，包括Twitter的高层管理人员。
Elon Musk：你好，你想讨论什么？
提问者1：你为什么要收购Twitter？
Elon Musk：Twitter已经是城市广场，影响力惊人。我相信我可以让它变得更好。这将使Twitter走上经济上可持续发展的道路，有助于保护文明的未来。我还为Twitter提出了许多令人兴奋的新的和有用的功能。这将会是一个很大的乐趣。
Elon Musk：你好，你想讨论什么？
提问者2：你是谁？
Elon Musk。我是Elon Musk，一个众所周知的科技企业家和许多公司的首席执行官。你呢，你是做什么的？
Elon Musk: 你好，你想讨论什么？
提问者3：你和neuralink的一位执行官生了双胞胎吗？
Elon Musk：我是在尽我所能帮助人口不足的危机。出生率的崩溃是迄今为止人类文明面临的最大危险。
提问者3： 你真的是Elon Musk吗？
Elon Musk: 是的! 我现在在德克萨斯州，在我们的特斯拉总部呆了一段时间。刚从加州回来，我们正在为SpaceX的发射做准备，明天即将发表另一个TED演讲，这次是关于Neuralink。
Elon Musk: 你好，你想讨论什么？
提问者4：你开的是什么类型的车？
Elon Musk: 特斯拉Model S最多。时常开一辆Model X。
提问者4：你多大了？
Elon Musk: 51岁，但我说我有40岁人的精力。在健康方面最重要的事情不仅仅是避免摄入过多的糖，还有高强度的运动。
""",
Task.REWRITE: """有这样一段文本，{医生微笑着递给小明棒棒糖，同时让小明服下了药。}
改写这段话让它变得更加惊悚。{医生眼露凶光让小明服药，小明感到非常害怕}。
有这样一段文本，{雨下得很大}
改写这段话让它变得更加具体。{一霎时，雨点连成了线，大雨就像天塌了似的铺天盖地从空中倾泻下来。}。
有这样一段文本，{王老师离开了电影院，外面已经天黑了}
改写这段话让它包含更多电影信息。{这部电影比小王预想的时间要长，虽然口碑很好，但离开电影院时，小王还是有些失望。}
有这样一段文本，{男人站在超市外面打电话}
改写这段话来描述小丑。{男人站在马戏团外一边拿着气球一边打电话}
有这样一段文本，{风铃声响起}
改写这段话写的更加丰富。{我对这个风铃的感情是由它的铃声引起的。每当风吹来时，风铃发出非常动听的声音，听起来是那么乐观、豁达，像一个小女孩格格的笑声。}
""",
Task.COMPLETION: """
""",
Task.FREE: ""
}
task2prompt_end = {
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.COPY: """商品描述:冬季百搭的外套
文案:""",
Task.QA: """问题：四大名著分别是？
回答：""",
Task.TRANSLATE: """“I am a programmer in Tencent”的中文翻译是：""",
Task.DIALOG_JOURNAL: """Elon Musk: 你好，你想讨论什么？
我：收购Twitter之后你想做什么？
Elon Musk：""",
Task.REWRITE: """有这样一段文本，{我想家了}
改写这段话包含更多悲伤的感情。{""",
Task.COMPLETION: """“八月十八潮，壮观天下无。”这是北宋大诗人苏东坡咏赞钱塘秋潮的千古名句。千百年来，钱塘江以其奇特卓绝的江潮，不知倾倒了多少游人看客。
每年的农历八月十八前后，是观潮的最佳时节。这期间，秋阳朗照，金风宜人，钱塘江口的海塘上，游客群集，兴致盎然，争睹奇景。""",
Task.FREE: ""
}
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)
all_input = prompt_fix + user_input
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=False):
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}出错: internal error')
else:
text = choice.get("text", "")
text = cut_message(text)
if len(text) == 0:
st.info(f'生成结果#{idx}: 结果为空，可能的原因：生成的第一个字符为stop字符，请合理配置prompt或stop')
else:
st.success(f'生成结果#{idx}: {text}')
if task_type == 'COMPLETION':
st.text('Tips: 可多次生成后复制你认为的最好结果拼接于原文后，让WeLM继续生成。')
except Exception as e:
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)
if st.button('立即生成'):
completion()
footer="""
"""
st.markdown(footer,unsafe_allow_html=True)