import gradio as gr

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

text_generation_zh = pipeline(task=Tasks.text_generation, model='baichuan-inc/baichuan-7B',model_revision='v1.0.2')
text_generation_zh._model_prepare = True

def text_generation(input_text):
    if input_text == "":
        return ""
    return text_generation_zh(input_text)['text']

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # 百川-7B 模型体验
    输入文本以查看生成结果。
    """)
    inp = gr.Textbox(label="输入prompt")
    submit = gr.Button("提交")
    out = gr.Textbox(label="续写结果")
    submit.click(text_generation, inp, out)

    gr.Markdown("## Text Examples")
    gr.Examples(
        ["登鹳雀楼->王之涣\n夜雨寄北->", "Hamlet->Shakespeare\nOne Hundred Years of Solitude->"],
        inp,
        out,
        text_generation,
        cache_examples=True,
    )

if __name__ == "__main__":
    demo.queue().launch()