30 lines
685 B
Python
30 lines
685 B
Python
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]
|
||
)
|