import gradio as gr
import requests

import json
import os

APIKEY = os.environ.get("APIKEY")
APISECRET = os.environ.get("APISECRET")

def predict(text, seed, out_seq_length, min_gen_length, sampling_strategy, 
    num_beams, length_penalty, no_repeat_ngram_size, 
    temperature, topk, topp):
    global APIKEY
    global APISECRET

    url = 'https://wudao.aminer.cn/os/api/api/v2/completions_130B'

    payload = json.dumps({
        "apikey": APIKEY,
        "apisecret": APISECRET ,
        "language": "zh-CN",
        "prompt": text,
        "length_penalty": length_penalty,
        "temperature": temperature,
        "top_k": topk,
        "top_p": topp,
        "min_gen_length": min_gen_length,
        "sampling_strategy": sampling_strategy,
        "num_beams": num_beams,
        "max_tokens": out_seq_length
    })

    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    return response.json()['result']['output']['raw']


if __name__ == "__main__":
    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # GLM-130B
            An Open Bilingual Pre-Trained Model
            """)

        with gr.Row():
            with gr.Column():
                model_input = gr.Textbox(lines=7, placeholder='Input something in English or Chinese', label='Input')
                with gr.Row():
                    gen = gr.Button("Generate")
                    clr = gr.Button("Clear")
            outputs = gr.Textbox(lines=7, label='Output')
            

        seed = gr.Slider(maximum=100000, value=1234, step=1, label='Seed')
        out_seq_length = gr.Slider(maximum=256, value=128, minimum=8, step=1, label='Output Sequence Length')
        min_gen_length = gr.Slider(maximum=64, value=0, step=1, label='Min Generate Length')
        sampling_strategy = gr.Radio(choices=['BeamSearchStrategy', 'BaseStrategy'], value='BeamSearchStrategy', label='Search Strategy')

        with gr.Tabs():
            with gr.TabItem("Beam Search Parameter"):
                # beam search
                num_beams = gr.Slider(maximum=4, value=1, minimum=1, step=1, label='Number of Beams')
                length_penalty = gr.Slider(maximum=1, value=0.8, minimum=0, label='Length Penalty')
                no_repeat_ngram_size = gr.Slider(maximum=5, value=3, minimum=1, step=1, label='No Repeat Ngram Size')
            with gr.TabItem("Base Search Parameter"):
                # base search
                temperature = gr.Slider(maximum=1, value=1, minimum=0, label='Temperature')
                topk = gr.Slider(maximum=8, value=1, minimum=0, step=1, label='Top K')
                topp = gr.Slider(maximum=1, value=0, minimum=0, label='Top P')
            
        inputs = [model_input, seed, out_seq_length, min_gen_length, sampling_strategy, num_beams, length_penalty, no_repeat_ngram_size, temperature, topk, topp]
        gen.click(fn=predict, inputs=inputs, outputs=outputs)
        clr.click(fn=lambda value: gr.update(value=""), inputs=clr, outputs=model_input)

    demo.launch()