search_entityrelationtask.cc

/*
  CoPyright (c) by respective owners including Yahoo!, Microsoft, and
  individual contributors. All rights reserved.  Released under a BSD (revised)
  license as described in the file LICENSE.
*/
#include "search_entityrelationtask.h"
#include "vw.h"
using namespace std;

#define R_NONE 10 // label for NONE relation
#define LABEL_SKIP 11 // label for SKIP

namespace EntityRelationTask { Search::search_task task = { "entity_relation", run, initialize, finish, nullptr, nullptr };  }


namespace EntityRelationTask
{
using namespace Search;
namespace CS = COST_SENSITIVE;

void update_example_indicies(bool audit, example* ec, uint64_t mult_amount, uint64_t plus_amount);

struct task_data
{ float relation_none_cost;
  float entity_cost;
  float relation_cost;
  float skip_cost;
  bool constraints;
  bool allow_skip;
  v_array<uint32_t> y_allowed_entity;
  v_array<uint32_t> y_allowed_relation;
  size_t search_order;
  example* ldf_entity;
  example* ldf_relation;
};


void initialize(Search::search& sch, size_t& /*num_actions*/, po::variables_map& vm)
{ vw& all = sch.get_vw_pointer_unsafe();
  task_data * my_task_data = new task_data();
  sch.set_task_data<task_data>(my_task_data);

  new_options(all, "Entity Relation Options")
  ("relation_cost", po::value<float>(&(my_task_data->relation_cost))->default_value(1.0), "Relation Cost")
  ("entity_cost", po::value<float>(&(my_task_data->entity_cost))->default_value(1.0), "Entity Cost")
  ("constraints", "Use Constraints")
  ("relation_none_cost", po::value<float>(&(my_task_data->relation_none_cost))->default_value(0.5), "None Relation Cost")
  ("skip_cost", po::value<float>(&(my_task_data->skip_cost))->default_value(0.01f), "Skip Cost (only used when search_order = skip")
  ("search_order", po::value<size_t>(&(my_task_data->search_order))->default_value(0), "Search Order 0: EntityFirst 1: Mix 2: Skip 3: EntityFirst(LDF)" );
  add_options(all);

  check_option<size_t>(my_task_data->search_order, all, vm, "search_order", false, size_equal,
                       "warning: you specified a different value for --search_order than the one loaded from regressor. proceeding with loaded value: ", "");
  check_option<float>(my_task_data->relation_cost, all, vm, "relation_cost", false, float_equal,
                      "warning: you specified a different value for --relation_cost than the one loaded from regressor. proceeding with loaded value: ", "");
  check_option<float>(my_task_data->entity_cost, all, vm, "entity_cost", false, float_equal,
                      "warning: you specified a different value for --entity_cost than the one loaded from regressor. proceeding with loaded value: ", "");
  check_option<float>(my_task_data->relation_none_cost, all, vm, "relation_none_cost", false, float_equal,
                      "warning: you specified a different value for --relation_none_cost than the one loaded from regressor. proceeding with loaded value: ", "");
  check_option<float>(my_task_data->skip_cost, all, vm, "skip_cost", false, float_equal,
                      "warning: you specified a different value for --skip_cost than the one loaded from regressor. proceeding with loaded value: ", "");
  check_option(my_task_data->constraints, all, vm, "constraints", false,
               "warning: you specified a different value for --constraints than the one loaded from regressor. proceeding with loaded value: ");

  // setup entity and relation labels
  // Entity label 1:E_Other 2:E_Peop 3:E_Org 4:E_Loc
  // Relation label 5:R_Live_in 6:R_OrgBased_in 7:R_Located_in 8:R_Work_For 9:R_Kill 10:R_None
  my_task_data->constraints = vm.count("constraints") > 0;

  for(int i=1; i<5; i++)
    my_task_data->y_allowed_entity.push_back(i);

  for(int i=5; i<11; i++)
    my_task_data->y_allowed_relation.push_back(i);

  my_task_data->allow_skip = false;

  if(my_task_data->search_order != 3 && my_task_data->search_order != 4 )
  { sch.set_options(0);
  }
  else
  { example* ldf_examples = VW::alloc_examples(sizeof(CS::label), 10);
    CS::wclass default_wclass = { 0., 0, 0., 0. };
    for (size_t a=0; a<10; a++)
    { ldf_examples[a].l.cs.costs.push_back(default_wclass);
    }
    my_task_data->ldf_entity = ldf_examples;
    my_task_data->ldf_relation = ldf_examples+4;
    sch.set_options(Search::IS_LDF);
  }

  sch.set_num_learners(2);
  if(my_task_data->search_order == 4)
    sch.set_num_learners(3);

}

void finish(Search::search& sch)
{ task_data * my_task_data = sch.get_task_data<task_data>();
  my_task_data->y_allowed_entity.delete_v();
  my_task_data->y_allowed_relation.delete_v();
  if(my_task_data->search_order == 3)
  { for (size_t a=0; a<10; a++)
      VW::dealloc_example(CS::cs_label.delete_label, my_task_data->ldf_entity[a]);
    free(my_task_data->ldf_entity);
  }
  delete my_task_data;
}    // if we had task data, we'd want to free it here

bool check_constraints(size_t ent1_id, size_t ent2_id, size_t rel_id)
{ size_t valid_ent1_id [] = {2,3,4,2,2}; // encode the valid entity-relation combinations
  size_t valid_ent2_id [] = {4,4,4,3,2};
  if(rel_id - 5 == 5)
    return true;
  if(valid_ent1_id[rel_id-5] == ent1_id && valid_ent2_id[rel_id-5] == ent2_id)
    return true;
  return false;
}

void decode_tag(v_array<char> tag, char& type, int& id1, int& id2)
{ string s1;
  string s2;
  type = tag[0];
  size_t idx = 2;
  while(idx < tag.size() && tag[idx] != '_' && tag[idx] != '\0')
  { s1.push_back(tag[idx]);
    idx++;
  }
  id1 = atoi(s1.c_str());
  idx++;
  assert(type == 'R');
  while(idx < tag.size() && tag[idx] != '_' && tag[idx] != '\0')
  { s2.push_back(tag[idx]);
    idx++;
  }
  id2 = atoi(s2.c_str());
}

size_t predict_entity(Search::search&sch, example* ex, v_array<size_t>& /*predictions*/, ptag my_tag, bool isLdf=false)
{

  task_data* my_task_data = sch.get_task_data<task_data>();
  size_t prediction;
  if(my_task_data->allow_skip)
  { v_array<uint32_t> star_labels = v_init<uint32_t>();
    star_labels.push_back(ex->l.multi.label);
    star_labels.push_back(LABEL_SKIP);
    my_task_data->y_allowed_entity.push_back(LABEL_SKIP);
    prediction = Search::predictor(sch, my_tag).set_input(*ex).set_oracle(star_labels).set_allowed(my_task_data->y_allowed_entity).set_learner_id(1).predict();
    my_task_data->y_allowed_entity.pop();
  }
  else
  { if(isLdf)
    { for(uint32_t a=0; a<4; a++)
      { VW::copy_example_data(false, &my_task_data->ldf_entity[a], ex);
        update_example_indicies(true, &my_task_data->ldf_entity[a], 28904713, 4832917 * (uint64_t)(a+1));
        CS::label& lab = my_task_data->ldf_entity[a].l.cs;
        lab.costs[0].x = 0.f;
        lab.costs[0].class_index = a;
        lab.costs[0].partial_prediction = 0.f;
        lab.costs[0].wap_value = 0.f;
      }
      prediction = Search::predictor(sch, my_tag).set_input(my_task_data->ldf_entity, 4).set_oracle(ex->l.multi.label-1).set_learner_id(1).predict() + 1;
    }
    else
    { prediction = Search::predictor(sch, my_tag).set_input(*ex).set_oracle(ex->l.multi.label).set_allowed(my_task_data->y_allowed_entity).set_learner_id(0).predict();
    }
  }

  // record loss
  float loss = 0.0;
  if(prediction == LABEL_SKIP)
  { loss = my_task_data->skip_cost;
  }
  else if(prediction !=  ex->l.multi.label)
    loss= my_task_data->entity_cost;
  sch.loss(loss);
  return prediction;
}
size_t predict_relation(Search::search&sch, example* ex, v_array<size_t>& predictions, ptag my_tag, bool isLdf=false)
{ char type;
  int id1, id2;
  task_data* my_task_data = sch.get_task_data<task_data>();
  size_t hist[2];
  decode_tag(ex->tag, type, id1, id2);
  v_array<uint32_t> constrained_relation_labels = v_init<uint32_t>();
  if(my_task_data->constraints && predictions[id1]!=0 &&predictions[id2]!=0)
  { hist[0] = predictions[id1];
    hist[1] = predictions[id2];
  }
  else
  { hist[0] = 0;
    hist[1] = 0;
  }
  for(size_t j=0; j< my_task_data->y_allowed_relation.size(); j++)
  { if(!my_task_data->constraints || hist[0] == (size_t)0  || check_constraints(hist[0], hist[1], my_task_data->y_allowed_relation[j]))
      constrained_relation_labels.push_back(my_task_data->y_allowed_relation[j]);

  }

  size_t prediction;
  if(my_task_data->allow_skip)
  { v_array<uint32_t> star_labels = v_init<uint32_t>();
    star_labels.push_back(ex->l.multi.label);
    star_labels.push_back(LABEL_SKIP);
    constrained_relation_labels.push_back(LABEL_SKIP);
    prediction = Search::predictor(sch, my_tag).set_input(*ex).set_oracle(star_labels).set_allowed(constrained_relation_labels).set_learner_id(2).add_condition(id1, 'a').add_condition(id2, 'b').predict();
    constrained_relation_labels.pop();
  }
  else
  { if(isLdf)
    { int correct_label = 0; // if correct label is not in the set, use the first one
      for(size_t a=0; a<constrained_relation_labels.size(); a++)
      { VW::copy_example_data(false, &my_task_data->ldf_relation[a], ex);
        update_example_indicies(true, &my_task_data->ldf_relation[a], 28904713, 4832917* (uint64_t)(constrained_relation_labels[a]));
        CS::label& lab = my_task_data->ldf_relation[a].l.cs;
        lab.costs[0].x = 0.f;
        lab.costs[0].class_index = constrained_relation_labels[a];
        lab.costs[0].partial_prediction = 0.f;
        lab.costs[0].wap_value = 0.f;
        if(constrained_relation_labels[a] == ex->l.multi.label)
        { correct_label = (int)a;
        }
      }
      size_t pred_pos = Search::predictor(sch, my_tag).set_input(my_task_data->ldf_relation, constrained_relation_labels.size()).set_oracle(correct_label).set_learner_id(2).predict();
      prediction = constrained_relation_labels[pred_pos];
    }
    else
    { prediction = Search::predictor(sch, my_tag).set_input(*ex).set_oracle(ex->l.multi.label).set_allowed(constrained_relation_labels).set_learner_id(1).predict();
    }
  }

  float loss = 0.0;
  if(prediction == LABEL_SKIP)
  { loss = my_task_data->skip_cost;
  }
  else if(prediction !=  ex->l.multi.label)
  { if(ex->l.multi.label == R_NONE)
    { loss = my_task_data->relation_none_cost;
    }
    else
    { loss= my_task_data->relation_cost;
    }
  }
  sch.loss(loss);
  constrained_relation_labels.delete_v();
  return prediction;
}

void entity_first_decoding(Search::search& sch, vector<example*> ec, v_array<size_t>& predictions, bool isLdf=false)
{ // ec.size = #entity + #entity*(#entity-1)/2
  size_t n_ent = (size_t)(sqrt(ec.size()*8+1)-1)/2;
  // Do entity recognition first
  for (size_t i=0; i<ec.size(); i++)
  { if(i< n_ent)
      predictions[i] = predict_entity(sch, ec[i], predictions, (ptag)i, isLdf);
    else
      predictions[i] = predict_relation(sch, ec[i], predictions, (ptag)i, isLdf);
  }
}

void er_mixed_decoding(Search::search& sch, vector<example*> ec, v_array<size_t>& predictions)
{ // ec.size = #entity + #entity*(#entity-1)/2
  uint32_t n_ent = (uint32_t)((sqrt(ec.size()*8+1)-1)/2);
  for(uint32_t t=0; t<ec.size(); t++)
  { // Do entity recognition first
    uint32_t count = 0;
    for (ptag i=0; i<n_ent; i++)
    { if(count ==t)
      { predictions[i] = predict_entity(sch, ec[i], predictions, i);
        break;
      }
      count++;
      for(uint32_t j=0; j<i; j++)
      { if(count ==t)
        { ptag rel_index = (ptag)(n_ent + (2*n_ent-j-1)*j/2 + i-j-1);
          predictions[rel_index] = predict_relation(sch, ec[rel_index], predictions, rel_index);
          break;
        }
        count++;
      }
    }
  }
}

void er_allow_skip_decoding(Search::search& sch, vector<example*> ec, v_array<size_t>& predictions)
{ task_data* my_task_data = sch.get_task_data<task_data>();
  // ec.size = #entity + #entity*(#entity-1)/2
  size_t n_ent = (size_t)(sqrt(ec.size()*8+1)-1)/2;

  bool must_predict = false;
  size_t n_predicts = 0;
  size_t p_n_predicts = 0;
  my_task_data->allow_skip = true;

  // loop until all the entity and relation types are predicted
  for(ptag t=0; ; t++)
  { ptag i = t % (uint32_t)ec.size();
    if(n_predicts == ec.size())
      break;

    if(predictions[i] == 0)
    { if(must_predict)
      { my_task_data->allow_skip = false;
      }
      size_t prediction = 0;
      if(i < n_ent)  // do entity recognition
      { prediction = predict_entity(sch, ec[i], predictions, i);
      }
      else     // do relation recognition
      { prediction = predict_relation(sch, ec[i], predictions, i);
      }

      if(prediction != LABEL_SKIP)
      { predictions[i] = prediction;
        n_predicts++;
      }

      if(must_predict)
      { my_task_data->allow_skip = true;
        must_predict = false;
      }
    }

    if(i == ec.size()-1)
    { if(n_predicts == p_n_predicts)
      { must_predict = true;
      }
      p_n_predicts = n_predicts;
    }
  }
}

void run(Search::search& sch, vector<example*>& ec)
{ task_data* my_task_data = sch.get_task_data<task_data>();

  v_array<size_t> predictions = v_init<size_t>();
  for(size_t i=0; i<ec.size(); i++)
  { predictions.push_back(0);
  }

  switch(my_task_data->search_order)
  { case 0:
      entity_first_decoding(sch, ec, predictions, false);
      break;
    case 1:
      er_mixed_decoding(sch, ec, predictions);
      break;
    case 2:
      er_allow_skip_decoding(sch, ec, predictions);
      break;
    case 3:
      entity_first_decoding(sch, ec, predictions, true); //LDF = true
      break;
    default:
      cerr << "search order " << my_task_data->search_order << "is undefined." << endl;
  }


  for(size_t i=0; i<ec.size(); i++)
  { if (sch.output().good())
      sch.output() << predictions[i] << ' ';
  }
  predictions.delete_v();
}
// this is totally bogus for the example -- you'd never actually do this!
void update_example_indicies(bool audit, example* ec, uint64_t mult_amount, uint64_t plus_amount)
{ for (features& fs : *ec)
    for (feature_index& idx : fs.indicies)
      idx = ((idx * mult_amount) + plus_amount);
}
}