import gradio as gr
from transformers import pipeline

#pipe = pipeline("text-generation")
#gr.Interface.from_pipeline(pipe).launch()

API_URL = "https://api-inference.huggingface.co/models/jslin09/bloom-560m-finetuned-fraud"
headers = {"Authorization": "Bearer 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"} # Read only

description = "Legal Document Drafting with BLOOM"
api_key="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"
examples=[
          ["闕很大明知金融帳戶之存摺、提款卡及密碼係供自己使用之重要理財工具，"],
          ["森上梅前明知其無資力支付酒店消費，亦無付款意願，竟意圖為自己不法之所有，"]
]

iface = gr.Interface.load(
    "huggingface/jslin09/bloom-560m-finetuned-fraud",
    title="Drafting",
    inputs=[
        gr.Textbox(lines=10, label="Prompt", value="闕很大明知金融帳戶之存摺、提款卡及密碼係供自己使用之重要理財工具，"),  # prompt
        gr.Slider(10, 200, step=10, value=100),  # token_count
        gr.Slider(0.2, 2.0, step=0.1, value=1.0),  # temperature
        gr.Slider(0.0, 1.0, step=0.05, value=0.8),  # top_p
        gr.Slider(0.0, 1.0, step=0.1, value=0.1),  # presencePenalty
        gr.Slider(0.0, 1.0, step=0.1, value=0.1),  # countPenalty
    ],
    outputs=gr.Textbox(label="生成的草稿", lines=28),
    description=description,
    examples=examples,
    api_key=api_key,
)

demo = gr.TabbedInterface(
    [iface], ["分頁標籤"],
    title="Legal Document Drafting",
)

demo.queue()
demo.launch(share=False)