ai-source-detector / convert_model.py
yaya36095's picture
Create convert_model.py
0a31b92 verified
raw
history blame contribute delete
684 Bytes
import torch
from transformers import AutoConfig
# ุชุญู…ูŠู„ ุงู„ู†ู…ูˆุฐุฌ ุงู„ุฃุตู„ูŠ
model_path = "best_model_improved.pth"
model = torch.load(model_path)
# ุฅู†ุดุงุก ู…ู„ู ุงู„ุชูƒูˆูŠู† ุฅุฐุง ู„ู… ูŠูƒู† ู…ูˆุฌูˆุฏุงู‹
config = AutoConfig.for_model("efficientnet-b0",
num_labels=2, # ุนุฏุฏ ุงู„ูุฆุงุช (ุญู‚ูŠู‚ูŠ/ู…ุฒูŠู)
label2id={"real": 0, "fake": 1},
id2label={0: "real", 1: "fake"})
config.save_pretrained("./")
# ุญูุธ ุงู„ู†ู…ูˆุฐุฌ ุจุตูŠุบุฉ pytorch_model.bin
torch.save(model.state_dict(), "pytorch_model.bin")
print("ุชู… ุชุญูˆูŠู„ ุงู„ู†ู…ูˆุฐุฌ ุจู†ุฌุงุญ!")