Inf Retriever V1 1.5b

🚀 INF-Retriever-v1-1.5B

INF-Retriever-v1-1.5B 是一款轻量级的基于大语言模型的稠密检索模型，专为中英文数据检索任务优化，在异构信息检索任务中表现卓越。

🚀 快速开始

INF-Retriever-v1-1.5B 是 INF-Retriever-v1 的轻量级版本，由 INF TECH 开发。它基于 gte-Qwen2-1.5B-instruct 模型构建，并针对检索任务进行了微调，尤其适用于中文和英文数据。

截至 2025 年 2 月 19 日，在参数少于 70 亿的模型中，INF-Retriever-v1-1.5B 在 AIR-Bench 24.04 和 24.05 版本的中英文双语子排行榜上均排名第一，展示了其在异构信息检索任务中的前沿性能。

✨ 主要特性

• 针对中英文检索优化：该模型使用了专注于检索的中英文数据集进行微调，显著提高了其在各种检索场景中的准确性和效率。
• 顶级性能表现：INF-Retriever-v1-1.5B 在 AIR-Bench 排行榜上取得了优异成绩，使其成为跨领域异构信息检索任务的首选模型。

📦 安装指南

文档中未提及安装步骤，故跳过此章节。

💻 使用示例

基础用法

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("infly/inf-retriever-v1-1.5b", trust_remote_code=True)
# In case you want to reduce the maximum length:
model.max_seq_length = 8192

queries = [
    "how much protein should a female eat",
    "summit define",
]
documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",
]

query_embeddings = model.encode(queries, prompt_name="query")
document_embeddings = model.encode(documents)

scores = (query_embeddings @ document_embeddings.T) * 100
print(scores.tolist())
# [[89.36092376708984, 69.16694641113281], [57.51953125, 79.65923309326172]]

高级用法

import torch
import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel


def last_token_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]


def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'Instruct: {task_description}\nQuery: {query}'


# Each query must come with a one-sentence instruction that describes the task
task = 'Given a web search query, retrieve relevant passages that answer the query'
queries = [
    get_detailed_instruct(task, 'how much protein should a female eat'),
    get_detailed_instruct(task, 'summit define')
]
# No need to add instruction for retrieval documents
documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments."
]
input_texts = queries + documents

tokenizer = AutoTokenizer.from_pretrained('infly/inf-retriever-v1-1.5b', trust_remote_code=True)
model = AutoModel.from_pretrained('infly/inf-retriever-v1-1.5b', trust_remote_code=True)

max_length = 8192

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt')
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())
# [[89.36091613769531, 69.16694641113281], [57.519447326660156, 79.65917205810547]]

📚 详细文档

🔧 技术细节

文档中未提及技术实现细节，故跳过此章节。

📄 许可证

该模型使用的许可证为 apache-2.0。

👥 贡献者

导师

魏楚 • 徐英辉 • 齐远

INF 记忆团队

杨俊涵 (junhanyang@inftech.ai) • 万家禾 • 姚逸晨 (eason.yyc@inftech.ai)

📖 引用

如果您发现我们的模型有用，请考虑引用：

@misc {infly-ai_2025,
	author       = { Junhan Yang, Jiahe Wan, Yichen Yao, Wei Chu, Yinghui Xu, Yuan Qi },
	title        = { inf-retriever-v1 (Revision 5f469d7) },
	year         = 2025,
	url          = { https://huggingface.co/infly/inf-retriever-v1 },
	doi          = { 10.57967/hf/4262 },
	publisher    = { Hugging Face }
}