import streamlit as st
import time
import requests

st.header("DeeplySorry")

all_input = st.text_area('模型输入', value="""
今天，我们正式发布名为 DeeplySorry 的大规模神经网络模型，它可以代替您向您珍惜的亲人、朋友、爱人道歉。\n""", height=100)



top_p = st.slider('top_p', 0.0, 1.0, 0.95)
temperature = st.slider('temperature', 0.0, 1.0, 0.85)
max_tokens = st.slider('max tokens', 4, 512, 64)
model_type = st.selectbox('model', ('large', 'xl'))


def completion(prompt):
    start = time.monotonic()
    resp = requests.post('https://welm.weixin.qq.com/v1/completions', json={
        'prompt': prompt,
        'model': model_type,
        'max_tokens': max_tokens,
        'temperature': temperature,
        'top_p': top_p,
        'n': 5,
        'stop': None,
        # 'stop': [[13, 13]],
    }, headers={"Authorization": f"Bearer {st.secrets['token']}"})
    if resp.status_code != 200:
        st.error(f'Bad response: {resp}, {resp.text}')
    else:
        answers = resp.json()
        st.json(answers)
        answers = [c['text'] for c in answers['choices'] if c['text'] is not None]
        cols = st.columns(3)
        for idx, answer in enumerate(answers):
            if idx >= 3:
                break
            with cols[idx]:
                content = (prompt + answer).replace("\n", "\n\n")
                st.markdown(f'## 版本{idx}\n\n{content}')
    end = time.monotonic()
    st.text(f'耗时：{end - start}')


if st.button('开始生成/换一批'):
    completion(all_input)