from pipeline import pipeline
from util import logger, print_func_and_docs


@print_func_and_docs
def chat_with_only_task():
    """只输入task为chat，有默认模型"""
    chat_pipeline = pipeline("chat")
    resp = chat_pipeline(messages=
        [
            {
                "role": "user",
                "content" : "hi"
            }
        ]
    )

@print_func_and_docs
def chat_with_only_model():
    """只输入模型"""
    chat_pipeline = pipeline(model="Baichuan/Baichuan2_7b_chat_pt")
    resp = chat_pipeline(messages=
        [
            {
                "role": "user",
                "content" : "hi"
            }
        ]
    )

@print_func_and_docs
def chat_with_task_and_model():
    """输入task和模型"""
    chat_pipeline = pipeline("chat", model="Baichuan/Baichuan2_7b_chat_pt")
    resp = chat_pipeline(messages=
        [
            {
                "role": "user",
                "content": "hi"
            }
        ]
    )

@print_func_and_docs
def chat_with_ms_framework():
    """ms框架"""
    chat_pipeline = pipeline("chat", model="Baichuan/Baichuan2_7b_chat_pt",framework="ms")
    resp = chat_pipeline(messages=
    [
        {
            "role": "user",
            "content": "hi"
        }
    ]
    )

@print_func_and_docs
def text2image_with_only_task():
    """只输入task为text2image，但是没有默认模型"""
    try:
        text2image_pipeline = pipeline("text2image")
        resp = text2image_pipeline(prompt = "beautiful flower, oil painting style")
    except Exception as e:
        logger.error(e)


@print_func_and_docs
def text2image_with_extra_param():
    """输入task和model， 额外输入width, height"""
    text2image_pipeline = pipeline("text2image", model="PyTorch-NPU/stable-diffusion-2-1")
    resp = text2image_pipeline(prompt = "beautiful flower, oil painting style", width=2000, height=1000)

@print_func_and_docs
def text2image_with_silicon():
    """输入task和model， 开启silicon 加速"""
    text2image_pipeline = pipeline("text2image", model="PyTorch-NPU/stable-diffusion-2-1", use_silicon=True)
    resp = text2image_pipeline(prompt = "beautiful flower, oil painting style", width=2000, height=1000)

@print_func_and_docs
def text2image_with_ms():
    """输入task和model， 使用ms framework"""
    try:
        text2image_pipeline = pipeline("text2image", model="PyTorch-NPU/stable-diffusion-2-1", framework="ms")
        resp = text2image_pipeline(prompt = "beautiful flower, oil painting style", width=2000, height=1000)
    except Exception as e:
        logger.error(e)

@print_func_and_docs
def common_pt_only_task():
    """未专门注册过的任务， 使用pt framework， 现在版本的pipeline可以都放在这里"""
    text2image_pipeline = pipeline("text_classification")
    resp = text2image_pipeline(text="候选文本")

@print_func_and_docs
def common_pt_only_model():
    """未专门注册过的任务， 使用pt framework， 现在版本的pipeline可以都放在这里"""
    text2image_pipeline = pipeline(model="PyTorch-NPU/resnet_50", framework="pt")
    resp = text2image_pipeline(pic= "http://123.png")

@print_func_and_docs
def common_ms():
    """未专门注册过的任务， 使用ms framework， 现在版本的pipeline可以都放在这里"""
    text2image_pipeline = pipeline(model="PyTorch-NPU/distilbert_base_uncased_finetuned_sst_2_english", framework="ms")
    resp = text2image_pipeline(text="Today is a good day")


def output_support_report():
    """输出当前openmind完整任务和支持模型，依赖列表"""
    pass

def output_docker_file():
    """指定模型，framework，生成容器镜像配置文件"""
    pass


if __name__ == '__main__':
    # chat_with_only_task()
    # chat_with_only_model()
    # chat_with_task_and_model()
    # chat_with_ms_framework()
    # text2image_with_only_task()
    # text2image_with_extra_param()
    # text2image_with_silicon()
    # text2image_with_ms()
    common_pt_only_task()
    # common_pt_only_model()
    # common_ms()
    # output_support_report()
    # output_docker_file()