original_docstring: 
          Forward function that calculate the score of a batch of triples.
          In the 'single' mode, sample is a batch of triple.
          In the 'head-batch' or 'tail-batch' mode, sample consists two part.
          The first part is usually the positive sample.
          And the second part is the entities in the negative samples.
          Because negative samples and positive samples usually share two elements 
          in their triple ((head, relation) or (relation, tail)).
          
docstring: Forward function that calculate the score of a batch of triples.
  In the 'single' mode, sample is a batch of triple.
  In the 'head-batch' or 'tail-batch' mode, sample consists two part.
  The first part is usually the positive sample.
  And the second part is the entities in the negative samples.
  Because negative samples and positive samples usually share two elements
  in their triple ((head, relation) or (relation, tail)).
docstring_tokens: Forward function that calculate the score of a batch of triples . In the ' single ' mode sample is a batch of triple . In the ' head - batch ' or ' tail - batch ' mode sample consists two part . The first part is usually the positive sample . And the second part is the entities in the negative samples . Because negative samples and positive samples usually share two elements in their triple (( head relation ) or ( relation tail )) .
code: def forward(self, sample, mode='single'):
          '''
          Forward function that calculate the score of a batch of triples.
          In the 'single' mode, sample is a batch of triple.
          In the 'head-batch' or 'tail-batch' mode, sample consists two part.
          The first part is usually the positive sample.
          And the second part is the entities in the negative samples.
          Because negative samples and positive samples usually share two elements 
          in their triple ((head, relation) or (relation, tail)).
          '''
  
          if mode == 'single':
              batch_size, negative_sample_size = sample.size(0), 1
              
              head = torch.index_select(
                  self.entity_embedding, 
                  dim=0, 
                  index=sample[:,0]
              ).unsqueeze(1)
              
              relation = torch.index_select(
                  self.relation_embedding, 
                  dim=0, 
                  index=sample[:,1]
              ).unsqueeze(1)
              
              tail = torch.index_select(
                  self.entity_embedding, 
                  dim=0, 
                  index=sample[:,2]
              ).unsqueeze(1)
              
          elif mode == 'head-batch':
              tail_part, head_part = sample
              batch_size, negative_sample_size = head_part.size(0), head_part.size(1)
              
              head = torch.index_select(
                  self.entity_embedding, 
                  dim=0, 
                  index=head_part.view(-1)
              ).view(batch_size, negative_sample_size, -1)
              
              relation = torch.index_select(
                  self.relation_embedding, 
                  dim=0, 
                  index=tail_part[:, 1]
              ).unsqueeze(1)
              
              tail = torch.index_select(
                  self.entity_embedding, 
                  dim=0, 
                  index=tail_part[:, 2]
              ).unsqueeze(1)
              
          elif mode == 'tail-batch':
              head_part, tail_part = sample
              batch_size, negative_sample_size = tail_part.size(0), tail_part.size(1)
              
              head = torch.index_select(
                  self.entity_embedding, 
                  dim=0, 
                  index=head_part[:, 0]
              ).unsqueeze(1)
              
              relation = torch.index_select(
                  self.relation_embedding,
                  dim=0,
                  index=head_part[:, 1]
              ).unsqueeze(1)
              
              tail = torch.index_select(
                  self.entity_embedding, 
                  dim=0, 
                  index=tail_part.view(-1)
              ).view(batch_size, negative_sample_size, -1)
              
          else:
              raise ValueError('mode %s not supported' % mode)
              
          model_func = {
              'TransE': self.TransE,
              'DistMult': self.DistMult,
              'ComplEx': self.ComplEx,
              'RotatE': self.RotatE,
              'pRotatE': self.pRotatE
          }
          
          if self.model_name in model_func:
              score = model_func[self.model_name](head, relation, tail, mode)
          else:
              raise ValueError('model %s not supported' % self.model_name)
          
          return score
code_tokens: def forward ( self , sample , mode = 'single' ) : if mode == 'single' : batch_size , negative_sample_size = sample . size ( 0 ) , 1 head = torch . index_select ( self . entity_embedding , dim = 0 , index = sample [ : , 0 ] ) . unsqueeze ( 1 ) relation = torch . index_select ( self . relation_embedding , dim = 0 , index = sample [ : , 1 ] ) . unsqueeze ( 1 ) tail = torch . index_select ( self . entity_embedding , dim = 0 , index = sample [ : , 2 ] ) . unsqueeze ( 1 ) elif mode == 'head-batch' : tail_part , head_part = sample batch_size , negative_sample_size = head_part . size ( 0 ) , head_part . size ( 1 ) head = torch . index_select ( self . entity_embedding , dim = 0 , index = head_part . view ( - 1 ) ) . view ( batch_size , negative_sample_size , - 1 ) relation = torch . index_select ( self . relation_embedding , dim = 0 , index = tail_part [ : , 1 ] ) . unsqueeze ( 1 ) tail = torch . index_select ( self . entity_embedding , dim = 0 , index = tail_part [ : , 2 ] ) . unsqueeze ( 1 ) elif mode == 'tail-batch' : head_part , tail_part = sample batch_size , negative_sample_size = tail_part . size ( 0 ) , tail_part . size ( 1 ) head = torch . index_select ( self . entity_embedding , dim = 0 , index = head_part [ : , 0 ] ) . unsqueeze ( 1 ) relation = torch . index_select ( self . relation_embedding , dim = 0 , index = head_part [ : , 1 ] ) . unsqueeze ( 1 ) tail = torch . index_select ( self . entity_embedding , dim = 0 , index = tail_part . view ( - 1 ) ) . view ( batch_size , negative_sample_size , - 1 ) else : raise ValueError ( 'mode %s not supported' % mode ) model_func = { 'TransE' : self . TransE , 'DistMult' : self . DistMult , 'ComplEx' : self . ComplEx , 'RotatE' : self . RotatE , 'pRotatE' : self . pRotatE } if self . model_name in model_func : score = model_func [ self . model_name ] ( head , relation , tail , mode ) else : raise ValueError ( 'model %s not supported' % self . model_name ) return score
short_docstring: Forward function that calculate the score of a batch of triples.
short_docstring_tokens: Forward function that calculate the score of a batch of triples .
comment: '''
          Forward function that calculate the score of a batch of triples.
          In the 'single' mode, sample is a batch of triple.
          In the 'head-batch' or 'tail-batch' mode, sample consists two part.
          The first part is usually the positive sample.
          And the second part is the entities in the negative samples.
          Because negative samples and positive samples usually share two elements 
          in their triple ((head, relation) or (relation, tail)).
          '''