[fix] update some codes for quesnet

Author: wintermelon008

EduNLP/Vector/quesnet/quesnet.py

@@ -6,7 +6,7 @@
 
 
 class QuesNetModel(Vector):
-    def __init__(self, pretrained_dir, img_dir=None, device="cpu", **kwargs):
+    def __init__(self, pretrained_dir, device="cpu", **kwargs):
         """
         Parameters
         ----------
@@ -18,7 +18,7 @@ def __init__(self, pretrained_dir, img_dir=None, device="cpu", **kwargs):
             image dir
         """
         self.device = torch.device(device)
-        self.model = QuesNet.from_pretrained(pretrained_dir, img_dir=img_dir).to(self.device)
+        self.model = QuesNet.from_pretrained(pretrained_dir).to(self.device)
         self.model.eval()
 
     def __call__(self, items: dict):

examples/pretrain/quesnet.ipynb

@@ -138,8 +138,15 @@
     "}\n",
     "\n",
     "# 当前仅支持linux下训练\n",
-    "# pretrain_quesnet(os.path.join(os.path.abspath(data_dir), 'quesnet_data.json'),\n",
-    "#                  output_dir, tokenizer, True, train_params)"
+    "pretrain_quesnet(\n",
+    "    path=os.path.join(os.path.abspath(data_dir),'quesnet_data.json'),\n",
+    "    output_dir=output_dir,\n",
+    "    tokenizer=tokenizer,\n",
+    "    img_dir=None,\n",
+    "    save_embs=True,\n",
+    "    load_embs=False,\n",
+    "    train_params=train_params\n",
+    ")"
    ]
   },
   {

tests/test_pretrain/test_pretrained_quesnet.py

@@ -68,10 +68,11 @@ def test_tokenizer(self, standard_luna_data, pretrained_tokenizer_dir):
     def test_train_quesnet(self, standard_luna_data, pretrained_model_dir):
         test_items = [
             {'ques_content': '有公式$\\FormFigureID{wrong1?}$和公式$\\FormFigureBase64{wrong2?}$，\
-                    如图$\\FigureID{000004d6-0479-11ec-829b-797d5eb43535}$,\
-                    若$x,y$满足约束条件$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$'},
+            如图$\\FigureID{000004d6-0479-11ec-829b-797d5eb43535}$,若$x,y$满足约束条件$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$'},
             {'ques_content': '如图$\\FigureID{000004d6-0479-11ec-829b-797d5eb43535}$, \
-                    若$x,y$满足约束条件$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$'}
+                若$x,y$满足约束条件$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$',
+             "ques_figure_ids": ["000004d6-0479-11ec-829b-797d5eb43535"],
+             "ques_figure_paths": ["../../static/test_data/quesnet_img/000004d6-0479-11ec-829b-797d5eb43535.png"]}
         ]
 
         ques_file = path_append(abs_current_dir(__file__),
@@ -139,7 +140,9 @@ def test_quesnet_t2v(self, pretrained_model_dir):
     def test_quesnet_i2v(self, pretrained_model_dir):
         items = [
             {'ques_content': '如图$\\FigureID{000004d6-0479-11ec-829b-797d5eb43535}$, \
-                若$x,y$满足约束条件$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$'}
+                若$x,y$满足约束条件$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$',
+             "ques_figure_ids": ["000004d6-0479-11ec-829b-797d5eb43535"],
+             "ques_figure_paths": ["../../static/test_data/quesnet_img/000004d6-0479-11ec-829b-797d5eb43535.png"]}
         ]
         img_dir = path_append(abs_current_dir(__file__),
                               "../../static/test_data/quesnet_img", to_str=True)