14B_lora_ze / handler.py
zemuwen's picture
Update handler.py
088c9cc verified
raw
history blame
576 Bytes
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]}