ChancesYuan commited on
Commit
57c3ee4
1 Parent(s): d1cac81

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -8,14 +8,14 @@ from higher.patch import monkeypatch as make_functional
8
  # from src.models.one_shot_learner import OneShotLearner
9
 
10
  ### load KGE model
11
- edit_origin_model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path="zjunlp/KGEditor", subfolder="E-FB15k237")
12
- edit_ex_model = EXBertForMaskedLM.from_pretrained(pretrained_model_name_or_path="zjunlp/KGEditor", subfolder="E-FB15k237")
13
 
14
  edit_learner = torch.load("./learner_checkpoint/edit/learner_params.pt", map_location=torch.device('cpu'))
15
  add_learner = torch.load("./learner_checkpoint/add/learner_params.pt", map_location=torch.device('cpu'))
16
 
17
- add_origin_model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path="zjunlp/KGEditor", subfolder="A-FB15k237")
18
- add_ex_model = EXBertForMaskedLM.from_pretrained(pretrained_model_name_or_path="zjunlp/KGEditor", subfolder="A-FB15k237")
19
 
20
  ### init inputs
21
  ent_name2id = defaultdict(str)
 
8
  # from src.models.one_shot_learner import OneShotLearner
9
 
10
  ### load KGE model
11
+ edit_origin_model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path="ChancesYuan/KGEditor_Edit_Test")
12
+ edit_ex_model = EXBertForMaskedLM.from_pretrained(pretrained_model_name_or_path="ChancesYuan/KGEditor_Edit_Test")
13
 
14
  edit_learner = torch.load("./learner_checkpoint/edit/learner_params.pt", map_location=torch.device('cpu'))
15
  add_learner = torch.load("./learner_checkpoint/add/learner_params.pt", map_location=torch.device('cpu'))
16
 
17
+ add_origin_model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path="ChancesYuan/KGEditor_Add_Test")
18
+ add_ex_model = EXBertForMaskedLM.from_pretrained(pretrained_model_name_or_path="ChancesYuan/KGEditor_Add_Test")
19
 
20
  ### init inputs
21
  ent_name2id = defaultdict(str)