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