from pydantic import BaseModel
from transformers import pipeline
from config import MODEL_ARGS

classifier = pipeline("zero-shot-classification", **MODEL_ARGS)


# 返回一个结构化的内容
class ClassifyResult(BaseModel):
    sequence: str
    rank: list
    scores: list
    prediction: str


def classify(text: str, labels: list):
    output = classifier(text, labels)

    # pprint(output)

    # 根据 score，寻找最高的 label
    prediction_rank = output['scores'].index(max(output['scores']))

    return ClassifyResult(
        sequence=text,
        rank=output['labels'],
        scores=output['scores'],
        prediction=output['labels'][prediction_rank]
    )