CodexTrouter / ProTalk_ModelExporter.py
prelington's picture
Create ProTalk_ModelExporter.py
d7b01dc verified
raw
history blame
253 Bytes
from transformers import AutoModelForCausalLM
from safetensors.torch import save_file
model_path = "microsoft/phi-2"
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto")
save_file(model.state_dict(), "ProTalkModel.safetensors")