import gradio as gr
import os
import sys
from generate import main

def generate_text(length, nsamples, prefix, temperature):
    my_prefix = "--prefix=" + prefix + ","
    sys.argv = [
        'generate.py', 
        f'--length={length}', 
        f'--nsamples={nsamples}', 
        my_prefix, 
        f'--temperature={temperature}', 
        '--model_path=model/model_epoch10'
    ]
    output = main()
    return output

input_length = gr.inputs.Number(label="生成文本长度", min_value=1, max_value=10000, default=1000)
input_nsamples = gr.inputs.Number(label="生成文本数量", min_value=1, max_value=10, default=3)
input_prefix = gr.inputs.Textbox(label="起始文本")
input_temperature = gr.inputs.Slider(label="生成温度", minimum=0.1, maximum=2.0, default=0.7, step=0.1)
output_text = gr.outputs.Textbox(label="生成的文本")

title = "GPT2中文文本生成器"
description = "输入起始文本、生成文本长度、生成文本数量和生成温度，生成中文文本。"

gr.Interface(fn=generate_text, inputs=[input_length, input_nsamples, input_prefix, input_temperature], outputs=output_text, title=title, description=description).launch()