diff --git a/test/temp_test.py b/test/temp_test.py index 24617e0..e2d3a9b 100755 --- a/test/temp_test.py +++ b/test/temp_test.py @@ -1,19 +1,42 @@ -# from transformers import PreTrainedModel -from transformers.modeling_utils import PreTrainedModel import torch -from torch import nn -from src_llm.configuration_dolly import DollyConfig -from src_llm.modeling_dolly import DollyModel -# model = DollyModel(DollyConfig()) -# model.save_pretrained("./dolly_model") -# model.config.save_pretrained("./dolly_model") +import math -config = DollyConfig.from_pretrained("./dolly_model") -model = DollyModel.from_pretrained("./dolly_model", config=config) +def llama_rope(x, seq_len, dim, device='cpu'): + """ + 实现 Llama 中的 Rotary Positional Encoding (RoPE) + + x: 输入的嵌入,大小为 (batch_size, seq_len, dim) + seq_len: 序列长度 + dim: 嵌入的维度 + device: 设备(默认为 'cpu') + """ + # 生成位置索引 + position = torch.arange(0, seq_len, dtype=torch.float32, device=device) + + # 计算每个维度的旋转角度 + freqs = torch.exp(torch.arange(0, dim, 2, dtype=torch.float32, device=device) * -(math.log(10000.0) / dim)) + + # 计算位置编码矩阵 + angle_rates = position.unsqueeze(1) * freqs.unsqueeze(0) + + # 使用 sin 和 cos 对每一对相邻维度进行旋转 + sin = torch.sin(angle_rates) + cos = torch.cos(angle_rates) + + # 将 sin 和 cos 分别填充到奇数和偶数维度 + encoding = torch.zeros((seq_len, dim), device=device) + encoding[:, 0::2] = sin + encoding[:, 1::2] = cos + + # 对输入的嵌入应用旋转 + x_rot = torch.matmul(x, encoding.unsqueeze(0).transpose(0, 1)) + + return x_rot -# input_ids = torch.tensor([[101, 1024, 102]]).long() -# outputs = model(input_ids) -# print(outputs) +batch_size = 1 +seq_len = 2 # 序列长度 +dim = 3 # 嵌入维度 +x = torch.randn(batch_size, seq_len, dim) -model_size = sum(t.numel() for t in model.parameters()) -print(f"Model Size: {model_size/1000**2:.1f}M parameters") \ No newline at end of file +encoded_x = llama_rope(x, seq_len, dim) +print(encoded_x.shape) # 输出:torch.Size([2, 10, 16]) \ No newline at end of file