Simplestories 30M

🚀 简单故事模型家族

简单故事（SimpleStories）模型是为可解释性研究而创建的小型模型家族，它基于简单故事数据集进行训练。

🚀 快速开始

💻 使用示例

基础用法

import torch
from transformers import AutoTokenizer, LlamaForCausalLM


MODEL_SIZE = "30M"
model_path = "SimpleStories/SimpleStories-{}".format(MODEL_SIZE)

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(model_path)
model.to("cuda")
model.eval()

prompt = "The curious cat looked at the"

inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
input_ids = inputs.input_ids.to("cuda")

eos_token_id = 1

with torch.no_grad():
    output_ids = model.generate(
        input_ids=input_ids,
        max_new_tokens=400,
        temperature=0.7,
        do_sample=True,
        eos_token_id=eos_token_id
)

output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(f"\nGenerated text:\n{output_text}")

📚 详细文档

🔍 模型变体

📊 性能比较

模型评估的生成质量指标：

📝 分词器

我们使用了一个自定义的 WordPiece 分词器，其词表大小仅为 4096。我们对数据集进行了词法分析和覆盖增益分析，以构建一个小型分词器，同时不影响生成质量。

📦 数据集

简单故事数据集是由先进的语言模型生成的短故事集合。它具有以下特点：

• 故事标注包含高级概念：主题、话题、风格等。
• 通过种子故事生成实现更高的语义和句法多样性。
• 由 2024 年的模型生成。
• 预先计算了多个 NLP 指标以辅助筛选。
• 英语数据集仅包含 ASCII 字符。

你可以在 arXiv 上阅读有关该数据集的论文。

🏋️‍ 训练

训练和评估脚本可在 https://github.com/danbraunai/simple_stories_train 访问。

📄 许可证

本项目采用 MIT 许可证。