工具

sentence_transformers.util 定义了不同的有用函数来处理文本嵌入。

辅助函数

sentence_transformers.util.community_detection(embeddings: Tensor | ndarray, threshold: float = 0.75, min_community_size: int = 10, batch_size: int = 1024, show_progress_bar: bool = False) list[list[int]][源代码][源代码]

快速社区检测功能。

在嵌入中找到所有社区,即距离接近(小于阈值)的嵌入。仅返回大于 min_community_size 的社区。社区按降序返回。每个列表的第一个元素是社区的中心点。

参数:
  • embeddings (torch.Tensor or numpy.ndarray) -- 输入嵌入。

  • threshold (float) -- 判断两个嵌入是否接近的阈值。默认为 0.75。

  • min_community_size (int) -- 被视为社区的最小规模。默认为 10。

  • batch_size (int) -- 用于计算余弦相似度分数的批量大小。默认为 1024。

  • show_progress_bar (bool) -- 是否在计算过程中显示进度条。默认为 False。

返回:

一个社区列表,每个社区表示为一个索引列表。

返回类型:

List[List[int]]

sentence_transformers.util.http_get(url: str, path: str) None[源代码][源代码]

将URL下载到磁盘上的指定路径。

参数:
  • url (str) -- 下载的URL。

  • path (str) -- 保存下载文件的路径。

抛出:

requests.HTTPError -- 如果HTTP请求返回非200状态码。

返回:

sentence_transformers.util.is_training_available() bool[源代码][源代码]

如果具备训练 Sentence Transformers 模型所需的依赖项,即 Huggingface 数据集和 Huggingface 加速,则返回 True。

sentence_transformers.util.mine_hard_negatives(dataset: Dataset, model: SentenceTransformer, corpus: list[str] | None = None, cross_encoder: CrossEncoder | None = None, range_min: int = 0, range_max: int | None = None, max_score: float | None = None, margin: float | None = None, num_negatives: int = 3, sampling_strategy: Literal['random', 'top'] = 'top', as_triplets: bool = True, batch_size: int = 32, faiss_batch_size: int = 16384, use_faiss: bool = False, verbose: bool = True) Dataset[源代码][源代码]

向一个包含 (锚点, 正样本) 对的数组中添加硬负样本,以创建 (锚点, 正样本, 负样本) 三元组或 (锚点, 正样本, 负样本_1, ..., 负样本_n) 元组。

硬负样本挖掘是一种通过添加硬负样本(即可能与锚文本相似但并非如此的文本)来提高数据集质量的技术。使用硬负样本可以提高基于该数据集训练的模型的性能。

此函数使用 SentenceTransformer 模型来嵌入数据集中的句子,然后在数据集中找到与每个锚定句子最接近的匹配项。接着,它从最接近的匹配项中采样负样本,并可选择使用 CrossEncoder 模型对候选者进行重新评分。

你可以通过多种方式影响候选负样本的选择:

  • range_min:最接近匹配项的最小排名,视为负样本:有助于跳过最相似的文本,以避免将实际为正样本的文本标记为负样本。

  • range_max: 作为负样本考虑的最接近匹配的最大排名:有助于限制从样本中抽取负样本的候选数量。较低的值使处理速度更快,但可能导致满足边际或最大分数条件的候选负样本较少。

  • max_score:视为负样本的最大分数:有助于跳过与锚点过于相似的候选者。

  • margin: 硬负样本挖掘的边距:用于跳过与锚点相似度在正样本对一定边距内的候选负样本。可以使用值0来强制负样本始终比正样本离锚点更远。

  • sampling_strategy: 负样本的采样策略:“top” 或 “random”。“top” 将始终将前 n 个候选样本作为负样本,而 “random” 将从满足边际或最大分数条件的候选样本中随机采样 n 个负样本。

示例

>>> from sentence_transformers.util import mine_hard_negatives
>>> from sentence_transformers import SentenceTransformer
>>> from datasets import load_dataset
>>> # Load a Sentence Transformer model
>>> model = SentenceTransformer("all-MiniLM-L6-v2")
>>>
>>> # Load a dataset to mine hard negatives from
>>> dataset = load_dataset("sentence-transformers/natural-questions", split="train")
>>> dataset
Dataset({
    features: ['query', 'answer'],
    num_rows: 100231
})
>>> dataset = mine_hard_negatives(
...     dataset=dataset,
...     model=model,
...     range_min=10,
...     range_max=50,
...     max_score=0.8,
...     margin=0.1,
...     num_negatives=5,
...     sampling_strategy="random",
...     batch_size=128,
...     use_faiss=True,
... )
Batches: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 784/784 [00:43<00:00, 17.83it/s]
Batches: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 784/784 [00:07<00:00, 99.60it/s]
Querying FAISS index: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 784/784 [00:00<00:00, 884.99it/s]
Metric       Positive       Negative     Difference
Count         100,231        431,255        431,255
Mean           0.6866         0.4289         0.2804
Median         0.7010         0.4193         0.2740
Std            0.1125         0.0754         0.0999
Min            0.0303         0.1720         0.1001
25%            0.6221         0.3747         0.1991
50%            0.7010         0.4193         0.2740
75%            0.7667         0.4751         0.3530
Max            0.9584         0.7743         0.7003
Skipped 1289492 potential negatives (25.23%) due to the margin of 0.1.
Skipped 39 potential negatives (0.00%) due to the maximum score of 0.8.
Could not find enough negatives for 69900 samples (13.95%). Consider adjusting the range_max, range_min, margin and max_score parameters if you'd like to find more valid negatives.
>>> # Note: The minimum similarity difference is 0.1001 due to our margin of 0.1
>>> dataset
Dataset({
    features: ['query', 'answer', 'negative'],
    num_rows: 431255
})
>>> dataset[0]
{
    'query': 'when did richmond last play in a preliminary final',
    'answer': "Richmond Football Club Richmond began 2017 with 5 straight wins, a feat it had not achieved since 1995. A series of close losses hampered the Tigers throughout the middle of the season, including a 5-point loss to the Western Bulldogs, 2-point loss to Fremantle, and a 3-point loss to the Giants. Richmond ended the season strongly with convincing victories over Fremantle and St Kilda in the final two rounds, elevating the club to 3rd on the ladder. Richmond's first final of the season against the Cats at the MCG attracted a record qualifying final crowd of 95,028; the Tigers won by 51 points. Having advanced to the first preliminary finals for the first time since 2001, Richmond defeated Greater Western Sydney by 36 points in front of a crowd of 94,258 to progress to the Grand Final against Adelaide, their first Grand Final appearance since 1982. The attendance was 100,021, the largest crowd to a grand final since 1986. The Crows led at quarter time and led by as many as 13, but the Tigers took over the game as it progressed and scored seven straight goals at one point. They eventually would win by 48 points – 16.12 (108) to Adelaide's 8.12 (60) – to end their 37-year flag drought.[22] Dustin Martin also became the first player to win a Premiership medal, the Brownlow Medal and the Norm Smith Medal in the same season, while Damien Hardwick was named AFL Coaches Association Coach of the Year. Richmond's jump from 13th to premiers also marked the biggest jump from one AFL season to the next.",
    'negative': "2018 NRL Grand Final The 2018 NRL Grand Final was the conclusive and premiership-deciding game of the 2018 National Rugby League season and was played on Sunday September 30 at Sydney's ANZ Stadium.[1] The match was contested between minor premiers the Sydney Roosters and defending premiers the Melbourne Storm. In front of a crowd of 82,688, Sydney won the match 21–6 to claim their 14th premiership title and their first since 2013. Roosters five-eighth Luke Keary was awarded the Clive Churchill Medal as the game's official man of the match."
}
>>> dataset.push_to_hub("natural-questions-hard-negatives", "triplet-all")
参数:
  • dataset (Dataset) -- 包含 (锚点, 正样本) 对的 数据集。

  • model (SentenceTransformer) -- 用于嵌入句子的 SentenceTransformer 模型。

  • corpus (List[str], optional) -- 包含文档的字符串列表,这些文档将作为候选负样本,除了 dataset 中的第二列之外。默认为 None,在这种情况下,dataset 中的第二列将专门用作负候选语料库。

  • cross_encoder (CrossEncoder, optional) -- 用于重新评分候选者的 CrossEncoder 模型。默认为 None。

  • range_min (int) -- 考虑为负样本的最接近匹配的最小秩。默认为 0。

  • range_max (int, optional) -- 最大排名,用于考虑作为负样本的最接近匹配项。默认为 None。

  • max_score (float, optional) -- 作为负面考虑的最大分数。默认为 None。

  • margin (float, optional) -- 硬负样本挖掘的边际。默认为 None。

  • num_negatives (int) -- 要采样的负样本数量。默认为 3。

  • sampling_strategy (Literal["random", "top"]) -- 负样本采样策略:“top” 或 “random”。默认为 “top”。

  • as_triplets (bool) -- 如果为 True,则为每个输入样本返回最多 num_negatives 个 (锚点, 正样本, 负样本) 三元组。如果为 False,则为每个输入样本返回 1 个 (锚点, 正样本, 负样本_1, ..., 负样本_n) 元组。默认为 True。

  • batch_size (int) -- 用于编码数据集的批量大小。默认为 32。

  • faiss_batch_size (int) -- FAISS top-k 搜索的批处理大小。默认为 16384。

  • use_faiss (bool) -- 是否使用 FAISS 进行相似性搜索。对于大型数据集,这可能是推荐的。默认为 False。

  • verbose (bool) -- 是否打印统计信息和日志。默认为 True。

返回:

包含 (锚点, 正样本, 负样本) 三元组或 (锚点, 正样本, 负样本_1, ..., 负样本_n) 元组的数据集。

返回类型:

Dataset

sentence_transformers.util.normalize_embeddings(embeddings: Tensor) Tensor[源代码][源代码]

归一化嵌入矩阵,使得每个句子嵌入具有单位长度。

参数:

embeddings (Tensor) -- 输入嵌入矩阵。

返回:

归一化的嵌入矩阵。

返回类型:

Tensor

sentence_transformers.util.paraphrase_mining(model, sentences: list[str], show_progress_bar: bool = False, batch_size: int = 32, query_chunk_size: int = 5000, corpus_chunk_size: int = 100000, max_pairs: int = 500000, top_k: int = 100, score_function: ~typing.Callable[[~torch.Tensor, ~torch.Tensor], ~torch.Tensor] = <function cos_sim>) list[list[float | int]][源代码][源代码]

给定一个句子/文本列表,此函数执行释义挖掘。它将所有句子与其他所有句子进行比较,并返回具有最高余弦相似度分数的成对列表。

参数:
  • model (SentenceTransformer) -- 用于嵌入计算的 SentenceTransformer 模型

  • sentences (List[str]) -- 字符串列表(文本或句子)

  • show_progress_bar (bool, optional) -- 绘制进度条。默认为 False。

  • batch_size (int, optional) -- 模型同时编码的文本数量。默认为32。

  • query_chunk_size (int, optional) -- 同时搜索与 #query_chunk_size 最相似的配对。减少此值以降低内存占用(会增加运行时间)。默认为 5000。

  • corpus_chunk_size (int, optional) -- 同时将一个句子与 #corpus_chunk_size 个其他句子进行比较。减少此值以降低内存占用(会增加运行时间)。默认为 100000。

  • max_pairs (int, optional) -- 返回的文本对的最大数量。默认为 500000。

  • top_k (int, optional) -- 对于每句话,我们最多检索 top_k 句其他句子。默认为 100。

  • score_function (Callable[[Tensor, Tensor], Tensor], optional) -- 计算分数的函数。默认使用余弦相似度。默认为 cos_sim。

返回:

返回一个三元组的列表,格式为 [score, id1, id2]

返回类型:

List[List[Union[float, int]]]

此函数在查询嵌入列表和语料库嵌入列表之间执行余弦相似度搜索。它可用于信息检索/语义搜索,适用于最多约100万条目的语料库。

参数:
  • query_embeddings (Tensor) -- 一个包含查询嵌入的二维张量。

  • corpus_embeddings (Tensor) -- 一个包含语料嵌入的二维张量。

  • query_chunk_size (int, optional) -- 同时处理100个查询。增加该值可以提高速度,但需要更多内存。默认为100。

  • corpus_chunk_size (int, optional) -- 每次扫描语料库 100k 条目。增加该值可以提高速度,但需要更多内存。默认为 500000。

  • top_k (int, optional) -- 检索前 k 个匹配条目。默认为 10。

  • score_function (Callable[[Tensor, Tensor], Tensor], optional) -- 计算分数的函数。默认情况下,使用余弦相似度。

返回:

每个查询对应一个条目的列表。每个条目是一个字典列表,包含键 'corpus_id' 和 'score',按余弦相似度分数降序排列。

返回类型:

List[List[Dict[str, Union[int, float]]]]

sentence_transformers.util.truncate_embeddings(embeddings: ndarray, truncate_dim: int | None) ndarray[源代码][源代码]
sentence_transformers.util.truncate_embeddings(embeddings: Tensor, truncate_dim: int | None) Tensor

截断嵌入矩阵。

参数:
  • embeddings (Union[np.ndarray, torch.Tensor]) -- 嵌入以截断。

  • truncate_dim (Optional[int]) -- 要截断句子嵌入的维度。None 表示不进行截断。

示例

>>> from sentence_transformers import SentenceTransformer
>>> from sentence_transformers.util import truncate_embeddings
>>> model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")
>>> embeddings = model.encode(["It's so nice outside!", "Today is a beautiful day.", "He drove to work earlier"])
>>> embeddings.shape
(3, 768)
>>> model.similarity(embeddings, embeddings)
tensor([[1.0000, 0.8100, 0.1426],
        [0.8100, 1.0000, 0.2121],
        [0.1426, 0.2121, 1.0000]])
>>> truncated_embeddings = truncate_embeddings(embeddings, 128)
>>> truncated_embeddings.shape
>>> model.similarity(truncated_embeddings, truncated_embeddings)
tensor([[1.0000, 0.8092, 0.1987],
        [0.8092, 1.0000, 0.2716],
        [0.1987, 0.2716, 1.0000]])
返回:

截断的嵌入。

返回类型:

Union[np.ndarray, torch.Tensor]

相似度度量

sentence_transformers.util.cos_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) Tensor[源代码][源代码]

计算两个张量之间的余弦相似度。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

矩阵,其中 res[i][j] = cos_sim(a[i], b[j])

返回类型:

Tensor

sentence_transformers.util.dot_score(a: list | ndarray | Tensor, b: list | ndarray | Tensor) Tensor[源代码][源代码]

计算所有 i 和 j 的点积 dot_prod(a[i], b[j])。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

矩阵,其中 res[i][j] = dot_prod(a[i], b[j])

返回类型:

Tensor

sentence_transformers.util.euclidean_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) Tensor[源代码][源代码]

计算两个张量之间的欧几里得相似度(即,负距离)。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

矩阵,其中 res[i][j] = -euclidean_distance(a[i], b[j])

返回类型:

Tensor

sentence_transformers.util.manhattan_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) Tensor[源代码][源代码]

计算两个张量之间的曼哈顿相似度(即负距离)。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

矩阵,其中 res[i][j] = -曼哈顿距离(a[i], b[j])

返回类型:

Tensor

sentence_transformers.util.pairwise_cos_sim(a: Tensor, b: Tensor) Tensor[源代码][源代码]

计算成对余弦相似度 cos_sim(a[i], b[i])。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

向量,其中 res[i] = cos_sim(a[i], b[i])

返回类型:

Tensor

sentence_transformers.util.pairwise_dot_score(a: Tensor, b: Tensor) Tensor[源代码][源代码]

计算成对点积 dot_prod(a[i], b[i])。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

向量,其中 res[i] = dot_prod(a[i], b[i])

返回类型:

Tensor

sentence_transformers.util.pairwise_euclidean_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor)[源代码][源代码]

计算张量对之间的欧几里得距离(即负距离)。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

向量,其中 res[i] = -euclidean_distance(a[i], b[i])

返回类型:

Tensor

sentence_transformers.util.pairwise_manhattan_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor)[源代码][源代码]

计算张量对之间的曼哈顿相似度(即负距离)。

参数:
  • a (Union[list, np.ndarray, Tensor]) -- 第一个张量。

  • b (Union[list, np.ndarray, Tensor]) -- 第二个张量。

返回:

向量,其中 res[i] = -曼哈顿距离(a[i], b[i])

返回类型:

Tensor