モデルマージ

merge-pytorch-model.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer, XLMRobertaModel
+
+# カスタムモデルの定義
+class CustomXLMRobertaModel(XLMRobertaModel):
+    def __init__(self, config):
+        super(CustomXLMRobertaModel, self).__init__(config)
+        self.sparse_linear = SparseLinear(config.hidden_size, 1)  # 適切な出力次元を設定
+
+    def forward(self, *args, **kwargs):
+        outputs = super(CustomXLMRobertaModel, self).forward(*args, **kwargs)
+        dense_embeddings = outputs.last_hidden_state
+        sparse_embeddings = self.sparse_linear(dense_embeddings)
+        return sparse_embeddings
+
+# カスタムレイヤーの定義
+class SparseLinear(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(SparseLinear, self).__init__()
+        self.linear = nn.Linear(input_dim, output_dim)
+
+    def forward(self, x):
+        return self.linear(x)
+
+# モデルとトークナイザーのロード
+model_name = "."  # ローカルディレクトリを指定
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+config = AutoModel.from_pretrained(model_name).config
+model = CustomXLMRobertaModel.from_pretrained(model_name, config=config)
+
+# カスタムレイヤーのインスタンスを作成
+input_dim = 1024  # Denseベクトルの次元（モデルのhidden_sizeに合わせる）
+output_dim = 1    # Sparseベクトルの次元（保存された重みに合わせる）
+sparse_linear = SparseLinear(input_dim, output_dim)
+
+# Sparse線形変換のロード
+sparse_linear_path = "sparse_linear.pt"
+sparse_linear_state_dict = torch.load(sparse_linear_path, weights_only=True)
+
+# state_dictのキーを変換
+sparse_linear_state_dict = {
+    f"linear.{key}": value for key, value in sparse_linear_state_dict.items()
+}
+
+# カスタムレイヤーにstate_dictをロード
+sparse_linear.load_state_dict(sparse_linear_state_dict)
+
+# カスタムレイヤーの重みを元のモデルの重みに追加
+model.sparse_linear.load_state_dict(sparse_linear.state_dict())