Merge branch 'master' of https://github.com/qibinc/KBRD

parlai/agents/ripplenet/ripplenet.py → parlai/agents/kbrd/kbrd.py

@@ -16,7 +16,7 @@
 from parlai.core.torch_agent import Output, TorchAgent
 from parlai.core.utils import round_sigfigs
 
-from .modules import RippleNet
+from .modules import KBRD
 
 
 
@@ -73,11 +73,11 @@ def nltk_tokenize(text):
 
     return full_text_embeddings
 
-class RipplenetAgent(TorchAgent):
+class KBRDAgent(TorchAgent):
     @classmethod
     def add_cmdline_args(cls, argparser):
         """Add command-line arguments specifically for this agent."""
-        super(RipplenetAgent, cls).add_cmdline_args(argparser)
+        super(KBRDAgent, cls).add_cmdline_args(argparser)
         agent = argparser.add_argument_group("Arguments")
         agent.add_argument("-ne", "--n-entity", type=int)
         agent.add_argument("-nr", "--n-relation", type=int)
@@ -94,14 +94,14 @@ def add_cmdline_args(cls, argparser):
             "-lr", "--learningrate", type=float, default=3e-3, help="learning rate"
         )
         agent.add_argument("-nb", "--num-bases", type=int, default=8)
-        RipplenetAgent.dictionary_class().add_cmdline_args(argparser)
+        KBRDAgent.dictionary_class().add_cmdline_args(argparser)
         return agent
 
     def __init__(self, opt, shared=None):
         super().__init__(opt, shared)
         init_model, is_finetune = self._get_init_model(opt, shared)
 
-        self.id = "RipplenetAgent"
+        self.id = "KBRDAgent"
         self.n_entity = opt["n_entity"]
         self.n_hop = opt["n_hop"]
         self.n_memory = opt["n_memory"]
@@ -125,7 +125,7 @@ def __init__(self, opt, shared=None):
             # entity_text_emb = _load_text_embeddings(entity2entityId, opt["dim"], abstract_path)
 
             # encoder captures the input text
-            self.model = RippleNet(
+            self.model = KBRD(
                 n_entity=opt["n_entity"],
                 n_relation=opt["n_relation"],
                 dim=opt["dim"],
@@ -154,9 +154,9 @@ def __init__(self, opt, shared=None):
                 opt["learningrate"],
             )
 
-        elif "ripplenet" in shared:
+        elif "kbrd" in shared:
             # copy initialized data from shared table
-            self.model = shared["ripplenet"]
+            self.model = shared["kbrd"]
             self.kg = shared["kg"]
             self.movie_ids = shared["movie_ids"]
             self.optimizer = shared["optimizer"]
@@ -204,7 +204,7 @@ def reset_metrics(self):
     def share(self):
         """Share internal states."""
         shared = super().share()
-        shared["ripplenet"] = self.model
+        shared["kbrd"] = self.model
         shared["kg"] = self.kg
         shared["movie_ids"] = self.movie_ids
         shared["optimizer"] = self.optimizer

parlai/agents/ripplenet/modules.py → parlai/agents/kbrd/modules.py

@@ -380,7 +380,7 @@ def _edge_list(kg, n_entity, hop):
 
     return [(h, t, relation_idx[r]) for h, t, r in edge_list if relation_cnt[r] > 1000], len(relation_idx)
 
-class RippleNet(nn.Module):
+class KBRD(nn.Module):
     def __init__(
         self,
         n_entity,
@@ -397,7 +397,7 @@ def __init__(
         entity_text_emb,
         num_bases
     ):
-        super(RippleNet, self).__init__()
+        super(KBRD, self).__init__()
 
         self.n_entity = n_entity
         self.n_relation = n_relation

parlai/agents/transformer_rec/modules.py

@@ -14,8 +14,8 @@
 
 from parlai.core.torch_generator_agent import TorchGeneratorModel
 from parlai.core.utils import neginf
-from parlai.agents.ripplenet.modules import RippleNet
-from parlai.agents.ripplenet.ripplenet import _load_kg_embeddings
+from parlai.agents.kbrd.modules import KBRD
+from parlai.agents.kbrd.kbrd import _load_kg_embeddings
 
 
 def _normalize(tensor, norm_layer):
@@ -565,7 +565,7 @@ def __init__(self, opt, dictionary):
             open(os.path.join(opt["datapath"], "redial", "entity2entityId.pkl"), "rb")
         )
         # entity_kg_emb = _load_kg_embeddings(entity2entityId, opt["dim"], "sub_joined_embeddings.tsv")
-        self.ripplenet = RippleNet(opt['n_entity'],opt['n_relation'],opt['dim'],opt['n_hop'],opt['kge_weight'],opt['l2_weight'],opt['n_memory'],opt['item_update_mode'],opt['using_all_hops'], kg, None, None, num_bases=8)
+        self.ripplenet = KBRD(opt['n_entity'],opt['n_relation'],opt['dim'],opt['n_hop'],opt['kge_weight'],opt['l2_weight'],opt['n_memory'],opt['item_update_mode'],opt['using_all_hops'], kg, None, None, num_bases=8)
         state_dict = torch.load('saved/both_rgcn_1')['model']
         # state_dict = OrderedDict([('ripplenet.' + key, state_dict[key]) for key in state_dict])
         self.ripplenet.load_state_dict(state_dict)
@@ -575,10 +575,10 @@ def __init__(self, opt, dictionary):
         # self.user_representation_to_bias_1 = nn.Linear(opt['dim'], 256)
         # self.user_representation_to_bias_2 = nn.Linear(256, 2048)
         # self.user_representation_to_bias_3 = nn.Linear(2048, len(dictionary))
-        for param in self.ripplenet.parameters():
+        for param in self.kbrd.parameters():
             param.requires_grad = False
 
-        # self.ripplenet.user_representation(item_list)
+        # self.kbrd.user_representation(item_list)
 
     def reorder_encoder_states(self, encoder_states, indices):
         enc, mask = encoder_states

parlai/core/torch_generator_agent.py

@@ -555,7 +555,7 @@ def train_step(self, batch):
         self.zero_grad()
 
         if getattr(batch, 'movies', None):
-            assert hasattr(self.model, 'ripplenet')
+            assert hasattr(self.model, 'kbrd')
             self.model.user_representation, _ = self.model.ripplenet.user_representation(batch.movies)
             self.model.user_representation = self.model.user_representation.detach()
         try:
@@ -594,7 +594,7 @@ def eval_step(self, batch):
         self.model.eval()
         cand_scores = None
         if getattr(batch, 'movies', None):
-            assert hasattr(self.model, 'ripplenet')
+            assert hasattr(self.model, 'kbrd')
             self.model.user_representation, _ = self.model.ripplenet.user_representation(batch.movies)
             self.model.user_representation = self.model.user_representation.detach()
 

parlai/tasks/redial/train_ripplenet.py → parlai/tasks/redial/train_kbrd.py

@@ -15,9 +15,9 @@
     parser.set_defaults(
         task="redial",
         dict_tokenizer="split",
-        model="ripplenet",
+        model="kbrd",
         dict_file="saved/tmp",
-        model_file="saved/ripplenet",
+        model_file="saved/kbrd",
         fp16=True,
         batchsize=256,
         n_entity=64368,

scripts/both.sh

@@ -3,6 +3,6 @@ let num_runs=32
 
 for i in $(seq 0 $((num_runs-1)));
 do
-    CUDA_VISIBLE_DEVICES=2 python parlai/tasks/redial/train_ripplenet.py -mf saved/both_rgcn_$i
+    CUDA_VISIBLE_DEVICES=2 python parlai/tasks/redial/train_kbrd.py -mf saved/both_rgcn_$i
 done
 

scripts/onlymovie.sh

@@ -3,6 +3,6 @@ let num_runs=32
 
 for i in $(seq 0 $((num_runs-1)));
 do
-    CUDA_VISIBLE_DEVICES=0 python parlai/tasks/redial/train_ripplenet.py -mf saved/onlymovie_$i
+    CUDA_VISIBLE_DEVICES=0 python parlai/tasks/redial/train_kbrd.py -mf saved/onlymovie_$i
 done
 

scripts/show_bias.py

@@ -42,7 +42,7 @@ def setup_args(parser=None):
     agent = create_agent(opt, requireModelExists=True)
     entity2entityId = pkl.load(open('data/redial/entity2entityId.pkl', 'rb'))
 
-    up, _ = agent.model.ripplenet.user_representation([
+    up, _ = agent.model.kbrd.user_representation([
         list(map(lambda x: entity2entityId[x], movie_entities))
     ])