8k_to_4k.py

8k_to_4k.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


input_path = "./model_in"
output_path = "./model_out"
max_shard_size = "5GB"
new_max_length = 4096

print("load model...start")
model = AutoModelForCausalLM.from_pretrained(input_path)
print("load model...done")
print(model.dtype)


print("load tokenizer...start")
tokenizer = AutoTokenizer.from_pretrained(input_path)
print("load tokenizer...done")

tokenizer.model_max_length = new_max_length

# New positional embedding
old_embeddings = model.transformer.wpe.weight
new_embeddings = old_embeddings[:new_max_length, :]

model.transformer.wpe.weight = torch.nn.Parameter(new_embeddings)

# Update Config
model.config.n_positions = new_max_length

# Save
model.save_pretrained(output_path, max_shard_size=max_shard_size)
tokenizer.save_pretrained(output_path)