coffee
17 lines
760 B
Python
Executable File
17 lines
760 B
Python
Executable File
from src_llm.configuration_dolly import DollyConfig
|
|
from transformers.modeling_rope_utils import dynamic_rope_update, ROPE_INIT_FUNCTIONS
|
|
import torch
|
|
|
|
device="cuda"
|
|
inv_freq = 1.0 / (10000 ** (torch.arange(0, 128, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / 128))
|
|
batch_size = 2
|
|
seq_len = 4
|
|
position_ids = torch.arange(seq_len).repeat(batch_size, 1).to(device=device)
|
|
inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(device)
|
|
position_ids_expanded = position_ids[:, None, :].float()
|
|
with torch.autocast(device_type="cuda"):
|
|
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|
cos = emb.cos()
|
|
sin = emb.sin()
|