I'm not an expert but there's https://jaxformer.com, https://github.com/karpathy/nanochat and, of course, the original "Attention is all you need" paper at https://arxiv.org/abs/1706.03762.