app.py

import gradio as gr
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM
import sys
from argparse import ArgumentParser

parser = ArgumentParser()
# parser.add_argument("model_name_or_path")
parser.add_argument("--port","-p", type=int, default=7860)
parser.add_argument("--device","-d",type=str, default="cpu")
args = parser.parse_args()

tokenizer = AutoTokenizer.from_pretrained("p208p2002/bloomz-zh-instruct-1b7")
model = AutoModelForCausalLM.from_pretrained("p208p2002/bloomz-zh-instruct-1b7")
generator = pipeline('text-generation', model=model,tokenizer=tokenizer,device=args.device)

PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with an input that provides further context. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Response:"
    ),
}

def on_submit(instruction,input=""):
    if input == "":
        prompt = PROMPT_DICT["prompt_no_input"].format_map({
            "instruction":instruction
        })
    
    else:
        prompt = PROMPT_DICT["prompt_input"].format_map({
            "instruction":instruction,
            "input":input
        })
        
    output = generator(
        prompt, 
        do_sample=True, 
        max_new_tokens=512,
        eos_token_id = tokenizer.eos_token_id,
        pad_token_id = tokenizer.pad_token_id,
        early_stopping=True,
    )
    generated_text:str = output[0]['generated_text']
    generated_text = generated_text.replace(prompt,"")
    generated_text = generated_text.strip().split("\n")[0]
        
    return generated_text

demo = gr.Interface(
    title=f"Instruction Model - {"p208p2002/bloomz-zh-instruct-1b7"}",
    article="Usage and License Notices: licensed for research use only.",
    fn=on_submit, 
    inputs=["text","text"], 
    outputs="text",
    examples=[
        ["Identify the odd one out.","Twitter, Instagram, Telegram"],
        ["How can we reduce air pollution?",""],
        ["Rewrite the given paragraph in a shorter, easier to understand form.","Although it is generally accepted that the internet has allowed us to connect with people all over the world, there are still those people who are not familiar with its basic functions, who don’t understand why it has become so commonplace, or what its true capabilities are."],
        ["Use the given data to calculate the median.","[2, 3, 7, 8, 10]"],
        ["What is the product of 6 and 2?",""],
        ["如何減少空氣污染?",""],
        ["世界第一高峰是?",""],
        ["臺灣第一高樓是?",""],
        ["以下哪個名子並不是臺灣人名","1.張美麗 2.東尼史塔克 3.王曉明"],
        ["如何製作蘋果派?給我詳細的步驟",""]
    ]
)

demo.launch(server_name="0.0.0.0",server_port=args.port)