import gradio as gr
from transformers import AutoModelForCausalLM
from transformers import BloomTokenizerFast
from transformers import pipeline, set_seed
import random

model_name = "bloom-560m"
model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud')
tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')

def rnd_generate(prompt):
    rnd_seed = random.randint(10, 500)
    set_seed(rnd_seed)
    inputs = tokenizer(prompt, return_tensors="pt") # 回傳的張量使用 Pytorch的格式。如果是 Tensorflow 格式的話，則指定為 "tf"。
    results = model.generate(inputs["input_ids"],
                       max_length=500,
                       num_return_sequences=1, # 產生 1 個句子回來。
                       do_sample=True,
                       temperature=0.75,
                       top_k=50, 
                       top_p=0.9)
    return tokenizer.decode(results[0])

def generate(prompt):
    result_length = len(prompt) + 4
    inputs = tokenizer(prompt, return_tensors="pt") # 回傳的張量使用 Pytorch的格式。如果是 Tensorflow 格式的話，則指定為 "tf"。
    results = model.generate(inputs["input_ids"],
                       num_return_sequences=2, # 產生 2 個句子回來。
                       max_length=result_length,
                       early_stopping=True,
                       do_sample=True, 
                       top_k=50, 
                       top_p=0.9
                      )
    return tokenizer.decode(results[0])

examples = [
    ["闕很大明知金融帳戶之存摺、提款卡及密碼係供自己使用之重要理財工具，"],
    ["梅友乾明知其無資力支付酒店消費，亦無付款意願，竟意圖為自己不法之所有，"],
    ["瓊道帕意圖為自己不法所有，基於竊盜之犯意，"]
]

prompts = [
    ["輸入寫書類的句子，讓電腦生成下一句。或是按以下的範例句子。"],
    ["輸入寫書類的開頭句子，讓電腦隨機生成整篇草稿。"]
]

with gr.Blocks() as demo:
    gr.Markdown(
    """
    <h1 style="text-align: center;">Legal Document Drafting</h1>
    """)
    with gr.Row() as row:
        with gr.Column():
            gr.Markdown("""
            <h3>摘要</h3>
            <p>
            隨著大型語言模型技術的發展，藉由微調預訓練的大型語言模型來解決自然語言處理的下游任務，已經是主流的範式。訓練法律專業領域的語言模型，需要有大量的法律文件，以便讓語言模型能學得法律術語以及法律文書格式的特殊性，因此，通常需要依賴大量人工標註的資料集進行訓練，而在法律領域的應用，取得大量人工標註的資料集是有實際上的困難，因此傳統的NLP方法應用在法律文件起草中的任務就受到了限制。本文實驗結果表明，以大量無標記的法律文件，在本地端電腦中微調大型預訓練語言模型，除了可以顯著提高微調後所得之模型在法律文件起草任務上的性能，為實現自動化法律文件起草提供了新的思路和方法，並同時保障了資訊隱私以及降低資訊安全等問題。
            </p>
            """)
        with gr.Column(scale=1, min_width=600):
            with gr.Tab("Writing Assist"):
                result = gr.components.Textbox(lines=7, label="Writing Assist", placeholder=prompts[0])
                prompt = gr.components.Textbox(lines=2, label="Prompt", placeholder=examples[0], visible=False)
                gr.Examples(examples, label='Examples', inputs=[prompt])
                prompt.change(generate, inputs=[prompt], outputs=[result])
                btn = gr.Button("Next sentence")
                btn.click(generate, inputs=[result], outputs=[result])
            with gr.Tab("Random Generative"):
                result2 = gr.components.Textbox(lines=7, label="Random Generative", show_label=True, placeholder=prompts[1])
                gr.Examples(examples, label='Examples', inputs=[result2])
                rnd_btn = gr.Button("Random Drafting")
                rnd_btn.click(rnd_generate, inputs=[result2], outputs=[result2])
    
if __name__ == "__main__":
    demo.launch()