Update app.py

app.py

@@ -1,38 +1,29 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM
+from transformers import BloomTokenizerFast
 from transformers import pipeline, set_seed
-import random
 
-generator = pipeline('text-generation', model='jslin09/bloom-560m-finetuned-fraud')
+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>')
 
-def rnd_generate(text):
-    rnd_seed = random.randint(10, 500)
-    set_seed(rnd_seed)
-    result = generator(text,
-                       max_length=500,
-                       num_return_sequences=1,
-                       do_sample=True,
-                       temperature=0.75,
+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=5, # 產生五個句子回來。
+                       max_length=result_length,
+                       early_stopping=True,
+                       do_sample=True, 
                        top_k=50, 
-                       top_p=0.9)
-    return result[0]["generated_text"]
-
-def generate(text):
-    set_seed(55)
-    result = generator(text,
-                       max_length=500,
-                       num_return_sequences=1,
-                       do_sample=True,
-                       temperature=0.75,
-                       top_k=50, 
-                       top_p=0.9)
-    return result[0]["generated_text"]
-
+                       top_p=0.9
+                      )
+    return tokenizer.decode(results[0])
 
 examples = [
     ["闕很大明知金融帳戶之存摺、提款卡及密碼係供自己使用之重要理財工具，"],
     ["梅友乾明知其無資力支付酒店消費，亦無付款意願，竟意圖為自己不法之所有，"],
-    ["王大明意圖為自己不法所有，基於竊盜之犯意，"],
-    ["陳小智知悉吳良醫院可配合假病患製作不實之診斷證明書、病歷資料，以供渠等向保險公司詐領住院保險給付，即意圖為自己不法之所有，"]
+    ["王大明意圖為自己不法所有，基於竊盜之犯意，"]
 ]
 
 with gr.Blocks() as demo:
@@ -42,17 +33,12 @@ with gr.Blocks() as demo:
     """)
     with gr.Row():
         with gr.Column():
-            prompt = gr.components.Textbox(lines=5, label="Input Prompt", placeholder=examples[0])
-            with gr.Row():
-                with gr.Column():
-                    btn = gr.Button("Random Drafting")
-                with gr.Column():
-                    btn2 = gr.Button("Drafting")
-            gr.Examples(examples, inputs=[prompt])
-        with gr.Column():    
-            result = gr.components.Textbox(lines=15, label="Generative")
-        btn.click(rnd_generate, inputs=[prompt], outputs=[result])
-        btn2.click(generate, inputs=[prompt], outputs=[result])
-
+            result = gr.components.Textbox(lines=7, label="生成的草稿", show_label=True)
+            prompt = gr.components.Textbox(lines=2, label="輸入提示文字", placeholder=examples[0],visible=False)
+            gr.Examples(examples, label='例句', inputs=[prompt])
+            prompt.change(generate, inputs=[prompt], outputs=[result])
+            btn = gr.Button("下一句")
+            btn.click(generate, inputs=[result], outputs=[result])
+    
 if __name__ == "__main__":
-    demo.launch() # 在遠端啟動時，需要 share=True 。
+    demo.launch()