from typing import Dict, List, Any | |
from transformers import pipeline | |
class EndpointHandler(): | |
def __init__(self, path=""): | |
# 初始化方法,加载模型 | |
self.pipeline = pipeline("text-classification", model=path) | |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
# 处理输入数据并返回预测结果 | |
inputs = data.get("inputs", "") | |
prediction = self.pipeline(inputs) | |
return {"response": conversation.generated_responses[0]} | |