Merge pull request #161 from nnnyt/dev

[REFACTOR] remove useless code & minor changes
bigdata-ustc · Mar 20, 2024 · 8a9b344 · 8a9b344
2 parents ddc432a + aefcc3b
commit 8a9b344
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 633 deletions.
diff --git a/EduNLP/ModelZoo/jiuzhang/__init__.py b/EduNLP/ModelZoo/jiuzhang/__init__.py
@@ -1,2 +1,2 @@
 from .jiuzhang import *
-from .modeling import CPTModel as JiuzhangModel
+from .modeling import CPTModel as Jiuzhang
diff --git a/EduNLP/ModelZoo/jiuzhang/jiuzhang.py b/EduNLP/ModelZoo/jiuzhang/jiuzhang.py
@@ -8,7 +8,7 @@
 from typing import List
 from ..rnn.harnn import HAM
 from transformers import BartConfig as JiuzhangConfig
-from .modeling import CPTModel as JiuzhangModel
+from .modeling import CPTModel as Jiuzhang
 
 
 __all__ = ["JiuzhangForPropertyPrediction", "JiuzhangForKnowledgePrediction"]
@@ -20,10 +20,10 @@ def __init__(self, pretrained_model_dir=None, head_dropout=0.5, init=True):
         jiuzhang_config = JiuzhangConfig.from_pretrained(pretrained_model_dir)
         if init:
             print(f'Load Jiuzhang from checkpoint: {pretrained_model_dir}')
-            self.jiuzhang = JiuzhangModel.from_pretrained(pretrained_model_dir, ignore_mismatched_sizes=True)
+            self.jiuzhang = Jiuzhang.from_pretrained(pretrained_model_dir, ignore_mismatched_sizes=True)
         else:
             print(f'Load Jiuzhang from config: {pretrained_model_dir}')
-            self.jiuzhang = JiuzhangModel(jiuzhang_config)
+            self.jiuzhang = Jiuzhang(jiuzhang_config)
         self.hidden_size = self.jiuzhang.config.hidden_size
         self.head_dropout = head_dropout
         self.dropout = nn.Dropout(head_dropout)
@@ -90,10 +90,10 @@ def __init__(self,
         jiuzhang_config = JiuzhangConfig.from_pretrained(pretrained_model_dir)
         if init:
             print(f'Load Jiuzhang from checkpoint: {pretrained_model_dir}')
-            self.jiuzhang = JiuzhangModel.from_pretrained(pretrained_model_dir, ignore_mismatched_sizes=True)
+            self.jiuzhang = Jiuzhang.from_pretrained(pretrained_model_dir, ignore_mismatched_sizes=True)
         else:
             print(f'Load Jiuzhang from config: {pretrained_model_dir}')
-            self.jiuzhang = JiuzhangModel(jiuzhang_config)
+            self.jiuzhang = Jiuzhang(jiuzhang_config)
         self.hidden_size = self.jiuzhang.config.hidden_size
         self.head_dropout = head_dropout
         self.dropout = nn.Dropout(head_dropout)