import numpy as np
from transformers import pipeline


# TODO 1.文本分类
def dm01():
    # 基于预训练模型构建模型对象
    model = pipeline(task='text-classification', model='../models/chinese_sentiment')
    # print("model-->", model)

    #  准备输入数据
    inputs = '我不太喜欢天安门'

    # 把数据送给模型，得到结果，model()->调用forward方法
    result = model(inputs)
    print("result-->", result)


# TODO 2. 特征提取任务
def dm02():
    # 得到模型对象
    model = pipeline(task='feature-extraction', model='../models/bert-base-chinese')
    # 把文本送给模型，得到结果
    result = model('人生该如何起头')
    print("result-->", result)
    # [1, 9, 768] -> 1句话，9个token，每个token使用一个768维度的向量表示
    # 1：一句话
    # 9:9个token，分别是：【CLS 人 生 该 如 何 起 头 SEP】
    # 768：维度数
    print('shape-->', np.array(result).shape)


# TODO 3.完形填空任务
def dm03():
    # 创建模型对象
    model = pipeline(task='fill-mask', model='../models/chinese-bert-wwm')
    # 把数据送给模型，得到结果
    # result = model("我明天想去[MASK]家吃饭")
    result = model("我想明天去[MASK]边钓鱼")
    print("result-->", result)


# TODO 4.阅读理解任务
def dm04():
    # 准备数据
    context = '我叫张三，我是一个程序员，我的喜好是打篮球。'
    questions = ['我是谁？', '我是做什么的？', '我的爱好是什么？']
    # 得到模型对象
    model = pipeline(task='question-answering', model='../models/chinese_pretrain_mrc_roberta_wwm_ext_large')
    # 把数据送给模型
    result = model(context=context, question=questions)
    print("result-->", result)


# TODO 5.文本摘要任务
def dm05():
    # 得到预训练模型对象
    model = pipeline(task='summarization', model='../models/distilbart-cnn-12-6')
    # 准备文本
    text = "BERT is a transformers model pretrained on a large corpus of English data " \
           "in a self-supervised fashion. This means it was pretrained on the raw texts " \
           "only, with no humans labelling them in any way (which is why it can use lots " \
           "of publicly available data) with an automatic process to generate inputs and " \
           "labels from those texts. More precisely, it was pretrained with two objectives:Masked " \
           "language modeling (MLM): taking a sentence, the model randomly masks 15% of the " \
           "words in the input then run the entire masked sentence through the model and has " \
           "to predict the masked words. This is different from traditional recurrent neural " \
           "networks (RNNs) that usually see the words one after the other, or from autoregressive " \
           "models like GPT which internally mask the future tokens. It allows the model to learn " \
           "a bidirectional representation of the sentence.Next sentence prediction (NSP): the models" \
           " concatenates two masked sentences as inputs during pretraining. Sometimes they correspond to " \
           "sentences that were next to each other in the original text, sometimes not. The model then " \
           "has to predict if the two sentences were following each other or not."
    xxx = """
    先帝创业未半而中道崩殂，今天下三分，益州疲弊，此诚危急存亡之秋也。然侍卫之臣不懈于内，忠志之士忘身于外者，盖追先帝之殊遇，欲报之于陛下也。
    诚宜开张圣听，以光先帝遗德，恢弘志士之气，不宜妄自菲薄，引喻失义，以塞忠谏之路也。
宫中府中，俱为一体，陟罚臧否，不宜异同。若有作奸犯科及为忠善者，宜付有司论其刑赏，以昭陛下平明之理，不宜偏私，使内外异法也。
侍中、侍郎郭攸之、费祎、董允等，此皆良实，志虑忠纯，是以先帝简拔以遗陛下。愚以为宫中之事，事无大小，悉以咨之，然后施行，必能裨补阙漏，有所广益。
将军向宠，性行淑均，晓畅军事，试用于昔日，先帝称之曰能，是以众议举宠为督。愚以为营中之事，悉以咨之，必能使行阵和睦，优劣得所。
亲贤臣，远小人，此先汉所以兴隆也；亲小人，远贤臣，此后汉所以倾颓也。
    """
    # 把文本送给模型，得到结果
    result = model(text)
    print("result-->", result)


def dm06():
    # 得到预训练模型对象
    model = pipeline(task='token-classification',
                     model='../models/roberta-base-finetuned-cluener2020-chinese')
    # 把数据送给模型
    result = model("我喜欢杭州西湖")
    print("result-->", result)


if __name__ == '__main__':
    # dm01()
    # dm02()
    # dm03()
    # dm04()
    # dm05()
    dm06()