license: mit language:

• ko
• en pipeline_tag: text-classification

在Amazon SageMaker上训练韩语Reranker模型

提供用于开发韩语Reranker的微调指南。

ko-reranker是基于BAAI/bge-reranker-larger针对韩语数据进行微调的模型。
更多详细信息请参考korean-reranker-git / AWS博客：使用韩语Reranker提升检索增强生成(RAG)性能

0. 特性

• 与嵌入模型不同，Reranker以问题和文档作为输入，直接输出相似度而非嵌入向量。

• 向Reranker输入问题和段落即可获得相关性评分。

• Reranker基于交叉熵损失优化，因此相关性评分不受特定范围限制。

1. 使用方法

• 使用Transformers

    def exp_normalize(x):
      b = x.max()
      y = np.exp(x - b)
      return y / y.sum()
    
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model.eval()

    pairs = [["我讨厌你", "我爱你"], \
             ["我喜欢你", "我对你的感情可能是爱"]]

    with torch.no_grad():
        inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
        scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
        scores = exp_normalize(scores.numpy())
        print (f'first: {scores[0]}, second: {scores[1]}')

• 使用SageMaker

import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel

try:
	role = sagemaker.get_execution_role()
except ValueError:
	iam = boto3.client('iam')
	role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

# Hub模型配置. https://huggingface.co/models
hub = {
	'HF_MODEL_ID':'Dongjin-kr/ko-reranker',
	'HF_TASK':'text-classification'
}

# 创建Hugging Face模型类
huggingface_model = HuggingFaceModel(
	transformers_version='4.28.1',
	pytorch_version='2.0.0',
	py_version='py310',
	env=hub,
	role=role, 
)

# 部署模型至SageMaker推理
predictor = huggingface_model.deploy(
	initial_instance_count=1, # 实例数量
	instance_type='ml.g5.large' # ec2实例类型
)

runtime_client = boto3.Session().client('sagemaker-runtime')
payload = json.dumps(
    {
        "inputs": [
            {"text": "我讨厌你", "text_pair": "我爱你"},
            {"text": "我喜欢你", "text_pair": "我对你的感情可能是爱"}
        ]
    }
)

response = runtime_client.invoke_endpoint(
    EndpointName="<endpoint-name>",
    ContentType="application/json",
    Accept="application/json",
    Body=payload
)

## 反序列化
out = json.loads(response['Body'].read().decode()) ## for json
print (f'Response: {out}')

2. 背景

• 上下文顺序影响准确率(Lost in Middle, Liu et al., 2023)

• 为何需要使用Reranker

  • 当前LLM并非输入越多上下文效果越好，相关度高的内容应排在前面才能获得更好答案
  • 语义搜索使用的相似度(相关度)评分不够精确（即排名靠前的就一定比靠后的更符合问题吗？）
    • 嵌入模型擅长捕捉文档背后的含义
    • 问题与正确答案在语义上未必相同(假设文档嵌入)
    • 近似最近邻算法(ANN)带来的偏差

3. Reranker模型

• [Cohere] Reranker

• [BAAI] bge-reranker-large

• [BAAI] bge-reranker-base

4. 数据集

• msmarco-triplets

  • 来自MS MARCO段落数据集的(问题，答案，负面)-三元组，499,184个样本
  • 该数据集为英文构成
  • 基于Amazon Translate进行翻译后使用

• 格式

{"query": str, "pos": List[str], "neg": List[str]}

• Query是问题，pos是正面文本列表，neg是负面文本列表。若问题没有负面文本，可从整个语料库随机抽取部分作为负面文本。

• 示例

{"query": "韩国的首都是？", "pos": ["美国的首都是华盛顿，日本是东京，韩国是首尔"], "neg": ["美国的首都是华盛顿，日本是东京，朝鲜是平壤"]}

5. 性能

• 评估集:

./dataset/evaluation/eval_dataset.csv

• 训练参数:

{
    "学习率": 5e-6,
    "fp16": True,
    "训练轮数": 3,
    "每设备训练批量大小": 1,
    "梯度累积步数": 32,
    "训练组大小": 3,
    "最大长度": 512,
    "权重衰减": 0.01,
}

6. 致谢

• 部分代码基于FlagEmbedding和KoSimCSE-SageMaker开发。

7. 引用

• 如果您觉得本仓库有用，请考虑点赞⭐并引用

8. 贡献者:

• 张东进博士 (AWS AI/ML解决方案架构师) | 邮件 | 领英 | Git |

9. 许可证

• FlagEmbedding采用MIT许可证。

10. 访问统计