from gradio.components import Textbox, Slider, Checkbox
import gradio as gr
from transformers import pipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

model = AutoModelForCausalLM.from_pretrained("ckip-joint/bloom-1b1-zh", use_cache=True)
tokenizer = AutoTokenizer.from_pretrained("ckip-joint/bloom-1b1-zh")

generator = pipeline('text-generation', model=model, tokenizer=tokenizer)

def generate(text, max_length=64, temperature=0.7, top_k=25, top_p=0.9, no_repeat_ngram_size=10, do_sample=True):
    result = generator(text,max_length=max_length,
                            temperature=temperature,
                            top_k=top_k,
                            top_p=top_p,
                            no_repeat_ngram_size=10,
                            do_sample=do_sample,
                            )
    return result[0]["generated_text"]

examples = [
    ["四月的某一天，天氣晴朗寒冷，",64,0.7,25,0.9,10,True],
    ["問：台灣最高的建築物是？答：",64,0.1,25,0.9,10,True],
]

demo = gr.Interface(
    fn=generate,
    inputs=[
      Textbox(lines=5, label="Input Text"),
      Slider(minimum=32, maximum=1024, value=64, label="Max Length"),
      Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05, label="Temperature"),
      Slider(minimum=1, maximum=99, value=25, step=5, label="Top k"),
      Slider(minimum=0.5, maximum=0.99, value=0.9, step=0.01, label="Top p"),
      Slider(minimum=1, maximum=999, value=10, step=1, label="No Repeat Ngram Size"),
      Checkbox(value=True, label="Do Sample"),
    ],
    outputs=Textbox(label="Generated Text"),
    examples=examples
)

demo.launch()