当前位置：首页 > news >正文

ragflow 通过HuggingFace 配置rerank模型

news 2025/8/21 13:09:45

本地启动召回模型

import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"from flask import Flask, request, jsonify
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as npapp = Flask(__name__)# 加载模型
model_name = "BAAI/bge-reranker-v2-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()
print("model ok")@app.route('/rerank', methods=['POST'])
def rerank():data = request.jsonprint(data)query = data['query']texts = data.get('texts', [])# 构建输入对pairs = [[query, text] for text in texts]print(f"构建了 {len(pairs)} 个查询-文档对")# 编码inputs = tokenizer(pairs, padding=True, truncation=True,return_tensors='pt', max_length=512)# 推理with torch.no_grad():scores = model(**inputs).logits.squeeze().tolist()# 格式化结果results = []for i, score in enumerate(scores):results.append({"index": i,"score": float(score)})return jsonify(results)if __name__ == '__main__':app.run(host='0.0.0.0', port=8080)print("api ok")

2 配置模型