如何用LitServe部署Hugging Face的BERT模型？

模型部署流程

1. 准备模型

首先确保已安装transformers库：

pip install transformers

2. 实现API类

from transformers import BertTokenizer, BertModel
import litserve as ls

class BertLitAPI(ls.LitAPI):
    def setup(self, device):
        # 加载预训练模型
        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
        self.model = BertModel.from_pretrained("bert-base-uncased").to(device)
        
    def decode_request(self, request):
        # 提取文本输入
        return request["text"]
        
    def predict(self, text):
        # 执行推理
        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
        outputs = self.model(**inputs)
        # 返回平均后的embedding
        return outputs.last_hidden_state.mean(dim=1).tolist()
        
    def encode_response(self, output):
        # 格式化输出
        return {"embedding": output}

3. 启动服务

server = ls.LitServer(BertLitAPI(), accelerator="auto")
server.run(port=8000)

4. 测试API

使用curl发送测试请求：

curl -X POST "http://127.0.0.1:8000/predict" -H "Content-Type: application/json" -d '{"text": "Hello, world!"}'

5. 性能调优建议

• 启用max_batch_size参数实现批量推理
• 使用torch.jit.trace对模型进行编译优化
• 对高频调用场景启用GPU内存预分配

本答案来源于文章《LitServe：快速部署企业级通用AI模型推理服务》