from typing import Dict, List, Any | |
from transformers import pipeline | |
class EndpointHandler(): | |
def __init__(self, path=""): | |
# 初始化对话模型 | |
self.pipeline = pipeline("conversational", model=path) | |
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: | |
# 从请求数据中获取输入文本 | |
inputs = data.get("inputs", "") | |
# 使用对话模型生成响应 | |
conversation = self.pipeline(inputs) | |
# 返回模型的响应 | |
return {"response": conversation.generated_responses[0]} | |