File size: 633 Bytes
3de8f53 c74efee 3de8f53 dfa198a 3de8f53 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
from typing import Dict, Any
from transformers import pipeline
class EndpointHandler:
def __init__(self, path=""):
# 初始化对话生成管道
self.pipeline = pipeline("text-generation", model=path)
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""
处理请求数据并返回对话生成结果
data args:
inputs (:obj: `str`)
Return:
A :obj:`dict`: 返回生成的对话结果
"""
inputs = data.get("inputs")
conversation = self.pipeline(inputs,max_length=4096)
return {"generated_text": conversation}
|