Transformer Demo

跟着原论文 "Attention Is All You Need" 的主要框架走，一步一步用 pytorch 搭建一个简单的 Transformer 示例。

本项目主要借鉴了 nlp-tutorial/5-1.Transformer ，并在其基础上进行模块文件拆分及重要代码注释，使其结构更加清晰。

• 2024.7.27 更新 :
  • 输入序列无需手动填充空位，更新utils/sen2vec()函数，自动将输入序列用填充字符填充至max_len大小
  • 更新utils/vec2sen()函数，自动去除末尾多余占位符，遇到第一个结束符后直接停止
  • 更新模型的greedy_decoder()，自动解码直到遇到全结束符或全占位符或长度超限
  • 使用 kv-cache 加快解码
• 2024.9.10 更新 :
  • 分支 moe中使用手写 MoE 模块代替原来的MLP

How to Start

python main.py

项目结构

• model:
  • config.json: 包含模型构建中需要用到的参数
  • vocab.json: 示例词表
• utils.py: 杂项函数
• pe.py: 位置编码模块
• mha.py: 多头注意力机制模块
• ffn.py: 前馈神经网络模块
• enc.py: 编码模块
• dec.py: 解码模块
• trm.py: Transformer 模型
• main.py: 主函数文件

原论文中关于模型搭建的重要图片

• 主要框架

• 位置编码

• 注意力机制

• 前馈神经网络