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 }
}