开通服务

百度智能云控制台 (baidu.com)

在平台搜索bce-reranker开通,价格 ¥0.002元/千tokens

image.png

获取AK,SK

百度智能云-管理中心 (baidu.com)
创建 Access Key ,复制它和Secret Key

本地挂脚本

脚本如下,注意替换自己的AK,SK

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import uvicorn
from fastapi import FastAPI, Security, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from typing import Optional, List
import qianfan
from qianfan.resources import Reranker

app = FastAPI()
security = HTTPBearer()
env_bearer_token = 'sk-xxx'

class QADocs(BaseModel):
    query: Optional[str]
    documents: Optional[List[str]]

class Singleton(type):
    def __call__(cls, *args, **kwargs):
        if not hasattr(cls, '_instance'):
            cls._instance = super().__call__(*args, **kwargs)
        return cls._instance

class ReRanker(metaclass=Singleton):
    def __init__(self):
        os.environ["QIANFAN_ACCESS_KEY"] = "你的AK"
        os.environ["QIANFAN_SECRET_KEY"] = "你的SK"
        self.reranker = Reranker(model="bce-reranker-base_v1")

    def compute_score(self, query: str, documents: List[str]):
        if len(documents) > 0:
            result = self.reranker.do(query, documents)
            return result.body['results']
        else:
            return None

class Chat(object):
    def __init__(self):
        self.reranker = ReRanker()

    def fit_query_answer_rerank(self, query_docs: QADocs) -> List:
        if query_docs is None or len(query_docs.documents) == 0:
            return []

        scores = self.reranker.compute_score(query_docs.query, query_docs.documents)

        new_docs = []
        for index, score in enumerate(scores):
            new_docs.append({"index": score['index'], "text": query_docs.documents[score['index']], "score": score['relevance_score']})
        results = [{"index": documents["index"], "relevance_score": documents["score"]} for documents in list(sorted(new_docs, key=lambda x: x["score"], reverse=True))]
        return results

@app.post('/v1/rerank')
async def handle_post_request(docs: QADocs, credentials: HTTPAuthorizationCredentials = Security(security)):
    token = credentials.credentials
    if env_bearer_token is not None and token != env_bearer_token:
        raise HTTPException(status_code=401, detail="Invalid token")
    chat = Chat()
    try:
        results = chat.fit_query_answer_rerank(docs)
        return {"results": results}
    except Exception as e:
        print(f"报错:\n{e}")
        return {"error": "重排出错"}

if __name__ == "__main__":
    token = os.getenv("ACCESS_TOKEN")
    if token is not None:
        env_bearer_token = token
    try:
        uvicorn.run(app, host='0.0.0.0', port=6006)
    except Exception as e:
        print(f"API启动失败!\n报错:\n{e}")

修改docker配置

docker-composer如下修改

"reRankModels": [
        {
            "model": "bce-reranker-base_v1",
            "name": "检索重排-base",
            "charsPointsPrice": 0,
            "requestUrl": "http://你的服务器ip:6006/v1/rerank",
            "requestAuth": "sk-xx"
        }
    ],
Last modification:July 6, 2024
如果觉得我的文章对你有用,请随意赞赏