add test file

test/temp_test.py

@@ -1,19 +1,42 @@
-# from transformers import PreTrainedModel
-from transformers.modeling_utils import PreTrainedModel
 import torch
-from torch import nn
+import math
-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")
 
-config = DollyConfig.from_pretrained("./dolly_model")
+def llama_rope(x, seq_len, dim, device='cpu'):
-model = DollyModel.from_pretrained("./dolly_model", config=config)
+    """
+    实现 Llama 中的 Rotary Positional Encoding (RoPE)
     
-# input_ids = torch.tensor([[101, 1024, 102]]).long() 
+    x: 输入的嵌入，大小为 (batch_size, seq_len, dim)
-# outputs = model(input_ids)
+    seq_len: 序列长度
-# print(outputs)
+    dim: 嵌入的维度
+    device: 设备（默认为 'cpu'）
+    """
+    # 生成位置索引
+    position = torch.arange(0, seq_len, dtype=torch.float32, device=device)
     
-model_size = sum(t.numel() for t in model.parameters())
+    # 计算每个维度的旋转角度
-print(f"Model Size: {model_size/1000**2:.1f}M parameters")
+    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
+
+batch_size = 1
+seq_len = 2  # 序列长度
+dim = 3  # 嵌入维度
+x = torch.randn(batch_size, seq_len, dim)
+
+encoded_x = llama_rope(x, seq_len, dim)
+print(encoded_x.shape)  # 输出：torch.Size([2, 10, 16])