Lettuce-512D-v2（移动端优化版）

这是一个移动优先的 sentence-transformers 模型，专为记忆系统和角色扮演 RAG 应用而设计。

它将句子和段落映射到 512 维的稠密向量空间。与标准模型不同，该模型经过精细改造，支持 4096 tokens 的上下文长度，并针对叙事流畅性和角色交互进行了专门微调。

模型详情

关键规格

• 模型类型： Sentence Transformer（蒸馏与微调）
• 基础架构： all-MiniLM-L6-v2（BERT）并进行了自定义改造。
• 最大序列长度： 4096 tokens（从 512 扩展）。
• 输出维度： 512（通过稠密层投影）。
• 大小： ~23 MB（Int8 量化 ONNX）。
• 延迟： CPU 上约 ~1.3ms（短文本）。
• 相似度函数： 余弦相似度。

为什么选择这个模型？

标准移动端模型（如 MiniLM）在两个方面表现不佳：上下文（限制为 512 tokens）和叙事逻辑（它们更倾向于关键词匹配而非故事氛围）。

Lettuce-v2 通过 4 步工程流程解决了这些问题：

1. 改造： 使用复制-重复初始化将绝对位置嵌入扩展到 4096。
2. 投影： 添加可训练的稠密层将 384d -> 512d 进行投影。
3. 蒸馏： 从 BAAI/bge-m3（最先进模型）进行知识蒸馏，以保留通用语义逻辑。
4. 微调： 在 Augmental（命运石之门） 数据集上训练，并混合 NLI 数据，以优先考虑叙事/角色扮演的语义对齐。

完整模型架构

SentenceTransformer(
  (0): Transformer({'max_seq_length': 4096, 'do_lower_case': False}) with Transformer model: BertModel 
  (1): Pooling({'word_embedding_dimension': 384, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, ...})
  (2): Dense({'in_features': 384, 'out_features': 512, 'bias': True, 'activation_function': 'torch.nn.modules.linear.Identity'})
)

使用方法

直接使用

首先安装 Sentence Transformers 库：

pip install -U sentence-transformers

然后你可以加载此模型并运行推理。

from sentence_transformers import SentenceTransformer

# 从 🤗 Hub 下载
model = SentenceTransformer("Zeolit/lettuce-emb-512d-v2")

# 对叙事/角色扮演文本运行推理
sentences = [
    '"You raised a flag."',
    '*I take a deep breath, my mind racing with possibilities...* "Leap to an even earlier time."',
    'The quick brown fox jumps over the lazy dog.'
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 512]

# 获取嵌入的相似度分数
similarities = model.similarity(embeddings, embeddings)

训练详情

训练数据集

该模型在混合数据集上训练，以平衡叙事风格与通用逻辑：

• 主要数据： Heralax/Augmental-Dataset（视觉小说 / 角色扮演日志）。
• 基础数据： sentence-transformers/all-nli（逻辑与推理锚点）。
• 总大小： 18,798 个训练样本。

训练日志

训练显示出平滑的收敛，在收益递减点（Epoch 2）准确停止，以防止过拟合。

训练超参数

• per_device_train_batch_size: 4（梯度累积模拟更大的批次）
• num_train_epochs: 2
• fp16: True
• loss: CosineSimilarityLoss

引用

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}