基于ab框架 & redis镜像的缓存服务构建

1. redis缓存服务构建

docker pull redis
docker run -itd -p 6379:6379 redis:latest
docker run -itd -p 6379:6379 -v /root/public/data:/data redis:latest

至此一个简易的redis服务已搭好 也可使用-v 挂载 通过下述代码获取redis缓存服务的ipadress

docker inspect redis:latest

可获取虚拟host ip, 如172.17.0.14

2. ab框架设置

2.1 config/config.py设置

示例如下

"""
基础配置
"""
MODEL_PATH = '/root/app/models'

# TESTING = True
# DEBUG = True

# 这个值只能由英文字母构成，不能包含任何符号！
# APP_NAME = 'simple'

HOST = '0.0.0.0'

PRINT_RESULT = False
LOG_LEVEL = 'DEBUG'

workers=1
REDIS = {
    'host': '172.17.0.14',
    'port': 6379
    #'password': '6379'
}

这里REDIS的设置至关重要，host/port参考redis缓存服务inspect获取 password取决于部署redis服务, config中requirepass是否设置，如若未设置，此步可省去password设置

2.2 algorithms/api.py设置

接下来需在具体路由函数处进行设置，示例如下：

@algorithm()
def encodes_texts(cache_client, texts, batch_size=64):
    all_vecs = []
    # 加载模型二郎神模型经过SimCSE优化后效果较好
    # 初始化模型
    key = '/'.join(texts)
    value = cache_client.bget(key)
    if value:
        return value
    from transformers import AutoTokenizer, AutoModelForMaskedLM

    model_p = os.path.join(model_dir, 'IDEA-CCNL/Erlangshen-SimCSE-110M-Chinese')
    token_p = os.path.join(model_dir, 'IDEA-CCNL/Erlangshen-SimCSE-110M-Chinese')
    model = AutoModelForMaskedLM.from_pretrained(model_p)
    tokenizer = AutoTokenizer.from_pretrained(token_p)
    net = torch.nn.DataParallel(model)
    model.to("cuda:0")
    def txt2emb(batch_txt):
        with torch.no_grad():
            inp = tokenizer(batch_txt, max_length=512, padding='max_length', truncation=True, return_tensors="pt")
            inp.to("cuda:0")
            out = net(**inp, output_hidden_states=True)
            res = out.hidden_states[-1][:, 0, :].cpu().numpy()

        return res.tolist()

    dataset = Dataset(texts, max_len=512)
    dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False)
    for batch in dataloader:
        vecs = txt2emb(batch)
        all_vecs.extend(vecs)
    cache_client.bset('/'.join(texts), all_vecs, ex=86400)
    return all_vecs

• 要使用缓存必须配置config.REDIS。配置好之后缓存客户端会自动生成cache_client参数 bget/bset接口 Since ab v2.6.3
• 上述代码中核心注意点：入参中的cache_client为缓存客户端，配置config.REDIS后，当pyab运行时，自动生成。
• 同前一点类似，postman构建入参时，可完全不必考虑cache_client的传入，实践显示，相同输入，redis缓存服务大幅度节省时间
• REF: https://github.com/aliyun/algorithm-base/blob/master/examples/cache/algorithms/demo.py