<|start_filename|>data_gen/indexes/yago_crosswikis_wiki.lua<|end_filename|>
-- Loads the merged p(e|m) index.
if not opt then
  cmd = torch.CmdLine()
  cmd:option('-root_data_dir', '', 'Root path of the data, $DATA_PATH.')
  cmd:text()
  opt = cmd:parse(arg or {})
  assert(opt.root_data_dir ~= '', 'Specify a valid root_data_dir path argument.')  
end


require 'torch'
tds = tds or require 'tds' 

dofile 'utils/utils.lua'
dofile 'entities/ent_name2id_freq/ent_name_id.lua'

ent_p_e_m_index = tds.Hash()

mention_lower_to_one_upper = tds.Hash()

mention_total_freq = tds.Hash()

local crosswikis_textfilename = opt.root_data_dir .. 'generated/crosswikis_wikipedia_p_e_m.txt'
print('==> Loading crosswikis_wikipedia from file ' .. crosswikis_textfilename)
local it, _ = io.open(crosswikis_textfilename)
local line = it:read()

local num_lines = 0
while (line) do
  num_lines = num_lines + 1
  if num_lines % 2000000 == 0 then
    print('Processed ' .. num_lines .. ' lines. ')
  end

  local parts = split(line , '\t')
  local mention = parts[1]

  local total = tonumber(parts[2])
  assert(total)
  if total >= 1 then
    ent_p_e_m_index[mention] = tds.Hash()
    mention_lower_to_one_upper[mention:lower()] = mention
    mention_total_freq[mention] = total
    local num_parts = table_len(parts)
    for i = 3, num_parts do
      local ent_str = split(parts[i], ',')
      local ent_wikiid = tonumber(ent_str[1])
      local freq = tonumber(ent_str[2])
      assert(ent_wikiid)
      assert(freq)
      ent_p_e_m_index[mention][ent_wikiid] = freq / (total + 0.0) -- not sorted
    end
  end
  line = it:read()
end

local yago_textfilename = opt.root_data_dir .. 'generated/yago_p_e_m.txt'
print('==> Loading yago index from file ' .. yago_textfilename)
it, _ = io.open(yago_textfilename)
line = it:read()

num_lines = 0
while (line) do
  num_lines = num_lines + 1
  if num_lines % 2000000 == 0 then
    print('Processed ' .. num_lines .. ' lines. ')
  end

  local parts = split(line , '\t')
  local mention = parts[1]

  local total = tonumber(parts[2])
  assert(total)
  if total >= 1 then
    mention_lower_to_one_upper[mention:lower()] = mention
    if not mention_total_freq[mention] then
      mention_total_freq[mention] = total
    else
      mention_total_freq[mention] = total + mention_total_freq[mention]
    end
    
    local yago_ment_ent_idx = tds.Hash()
    local num_parts = table_len(parts)
    for i = 3, num_parts do
      local ent_str = split(parts[i], ',')
      local ent_wikiid = tonumber(ent_str[1])
      local freq = 1
      assert(ent_wikiid)
      yago_ment_ent_idx[ent_wikiid] = freq / (total + 0.0) -- not sorted
    end
  
    if not ent_p_e_m_index[mention] then
      ent_p_e_m_index[mention] = yago_ment_ent_idx
    else
      for ent_wikiid,prob in pairs(yago_ment_ent_idx) do
        if not ent_p_e_m_index[mention][ent_wikiid] then
          ent_p_e_m_index[mention][ent_wikiid] = 0.0
        end
        ent_p_e_m_index[mention][ent_wikiid] = math.min(1.0, ent_p_e_m_index[mention][ent_wikiid] + prob)
      end
      
    end
  
  end
  line = it:read()
end

assert(ent_p_e_m_index['<NAME>'] and ent_p_e_m_index['<NAME>'])

-- Function used to preprocess a given mention such that it has higher 
-- chance to have at least one valid entry in the p(e|m) index.
function preprocess_mention(m)
  assert(ent_p_e_m_index and mention_total_freq)
  local cur_m = modify_uppercase_phrase(m)
  if (not ent_p_e_m_index[cur_m]) then
    cur_m = m
  end
  if (mention_total_freq[m] and (mention_total_freq[m] > mention_total_freq[cur_m])) then
    cur_m = m -- Cases like 'U.S.' are handed badly by modify_uppercase_phrase
  end
  -- If we cannot find the exact mention in our index, we try our luck to find it in a case insensitive index.
  if (not ent_p_e_m_index[cur_m]) and mention_lower_to_one_upper[cur_m:lower()] then
    cur_m = mention_lower_to_one_upper[cur_m:lower()]
  end
  return cur_m
end


print('    Done loading index')


<|start_filename|>entities/pretrained_e2v/e2v.lua<|end_filename|>
-- Loads pre-trained entity vectors trained using the file entity/learn_e2v/learn_a.lua

assert(opt.ent_vecs_filename)
print('==> Loading pre-trained entity vectors: e2v from file ' .. opt.ent_vecs_filename)

assert(opt.entities == 'RLTD', 'Only RLTD entities are currently supported. ALL entities would blow the GPU memory.')

-- Defining variables:
ent_vecs_size = 300

geom_w2v_M = w2vutils.M:float()

e2vutils = {}

-- Lookup table: ids -> tensor of vecs
e2vutils.lookup = torch.load(opt.root_data_dir .. 'generated/ent_vecs/' .. opt.ent_vecs_filename)
e2vutils.lookup = nn.Normalize(2):forward(e2vutils.lookup) -- Needs to be normalized to have norm 1.

assert(e2vutils.lookup:size(1) == get_total_num_ents() and
  e2vutils.lookup:size(2) == ent_vecs_size, e2vutils.lookup:size(1) .. ' ' .. get_total_num_ents())
assert(e2vutils.lookup[unk_ent_thid]:norm() == 0, e2vutils.lookup[unk_ent_thid]:norm())

-- ent wikiid -> vec
e2vutils.entwikiid2vec = function(self, ent_wikiid)
  local thid = get_thid(ent_wikiid)
  return self.lookup[thid]:float()
end
assert(torch.norm(e2vutils:entwikiid2vec(unk_ent_wikiid)) == 0)


e2vutils.entname2vec = function (self,ent_name)
  assert(ent_name)
  return e2vutils:entwikiid2vec(get_ent_wikiid_from_name(ent_name))
end

-- Entity similarity based on cosine distance (note that entity vectors are normalized).
function entity_similarity(e1_wikiid, e2_wikiid)
  local e1_vec = e2vutils:entwikiid2vec(e1_wikiid)
  local e2_vec = e2vutils:entwikiid2vec(e2_wikiid)
  return e1_vec * e2_vec
end


-----------------------------------------------------------------------
---- Some unit tests to understand the quality of these embeddings ----
-----------------------------------------------------------------------
local function geom_top_k_closest_words(ent_name, ent_vec, k)
  local tf_map = ent_wiki_words_4EX[ent_name]
  local w_not_found = {}
  for w,_ in pairs(tf_map) do
    if tf_map[w] >= 10 then
      w_not_found[w] = tf_map[w]
    end
  end
  
  distances = geom_w2v_M * ent_vec

  local best_scores, best_word_ids = topk(distances, k)
  local returnwords = {}
  local returndistances = {}
  for i = 1,k do
    local w = get_word_from_id(best_word_ids[i])
    if get_w_id_freq(best_word_ids[i]) >= 200 then
      local w_freq_str = '[fr=' .. get_w_id_freq(best_word_ids[i]) .. ']'
      if is_stop_word_or_number(w) then
        table.insert(returnwords, red(w .. w_freq_str))
      elseif tf_map[w] then
        if tf_map[w] >= 15 then
          table.insert(returnwords, yellow(w .. w_freq_str .. '{tf=' .. tf_map[w] .. '}'))
        else
          table.insert(returnwords, skyblue(w .. w_freq_str .. '{tf=' .. tf_map[w] .. '}'))
        end
        w_not_found[w] = nil
      else
        table.insert(returnwords, w .. w_freq_str)
      end
      assert(best_scores[i] == distances[best_word_ids[i]], best_scores[i] .. '  ' .. distances[best_word_ids[i]])
      table.insert(returndistances, distances[best_word_ids[i]])
    end
  end
  return returnwords, returndistances, w_not_found
end


local function geom_most_similar_words_to_ent(ent_name, k)
  local ent_wikiid = get_ent_wikiid_from_name(ent_name)
  local k = k or 1
  local ent_vec = e2vutils:entname2vec(ent_name)
  assert(math.abs(1 - ent_vec:norm()) < 0.01 or ent_vec:norm() == 0, ':::: ' .. ent_vec:norm())
  
  print('\nTo entity: ' .. blue(ent_name) .. '; vec norm = ' .. ent_vec:norm() .. ':')
  neighbors, scores, w_not_found = geom_top_k_closest_words(ent_name, ent_vec, k)
  print(green('TOP CLOSEST WORDS: ') .. list_with_scores_to_str(neighbors, scores))
  
  local str = yellow('WORDS NOT FOUND: ')
  for w,tf in pairs(w_not_found) do
    if tf >= 20 then
      str = str .. yellow(w .. '{' .. tf .. '}; ')
    else
      str = str .. w .. '{' .. tf .. '}; '
    end
  end
  print('\n' .. str)
  print('============================================================================')
end

function geom_unit_tests()
  print('\n' .. yellow('TOP CLOSEST WORDS to a given entity based on cosine distance:'))  
  print('For each word, we show the unigram frequency [fr] and the cosine similarity.')
  print('Infrequent words [fr < 500] have noisy embeddings, thus should be trusted less.')
  print('WORDS NOT FOUND contains frequent words from the Wikipedia canonical page that are not found in the TOP CLOSEST WORDS list.')
  
  for i=1,table_len(ent_names_4EX) do
    geom_most_similar_words_to_ent(ent_names_4EX[i], 300)
  end
end

print('    Done reading e2v data. Entity vocab size = ' .. e2vutils.lookup:size(1))


<|start_filename|>data_gen/gen_p_e_m/merge_crosswikis_wiki.lua<|end_filename|>
-- Merge Wikipedia and Crosswikis p(e|m) indexes
-- Run: th data_gen/gen_p_e_m/merge_crosswikis_wiki.lua -root_data_dir $DATA_PATH

cmd = torch.CmdLine()
cmd:option('-root_data_dir', '', 'Root path of the data, $DATA_PATH.')
cmd:text()
opt = cmd:parse(arg or {})
assert(opt.root_data_dir ~= '', 'Specify a valid root_data_dir path argument.')


require 'torch'
dofile 'utils/utils.lua'
dofile 'entities/ent_name2id_freq/ent_name_id.lua'

print('\nMerging Wikipedia and Crosswikis p_e_m')

tds = tds or require 'tds' 
merged_e_m_counts = tds.Hash()

print('Process Wikipedia')
it, _ = io.open(opt.root_data_dir .. 'generated/wikipedia_p_e_m.txt')
line = it:read()

while (line) do
  local parts = split(line, "\t")
  local mention  = parts[1]

  if (not mention:find('Wikipedia')) and (not mention:find('wikipedia')) then
    if not merged_e_m_counts[mention] then
      merged_e_m_counts[mention] = tds.Hash()
    end  

    local total_freq = tonumber(parts[2])
    assert(total_freq)
    local num_ents = table_len(parts)
    for i = 3, num_ents do
      local ent_str = split(parts[i], ",")
      local ent_wikiid = tonumber(ent_str[1])
      assert(ent_wikiid)
      local freq = tonumber(ent_str[2])
      assert(freq)

      if not merged_e_m_counts[mention][ent_wikiid] then
        merged_e_m_counts[mention][ent_wikiid] = 0
      end
      merged_e_m_counts[mention][ent_wikiid] = merged_e_m_counts[mention][ent_wikiid] + freq
    end
  end
  line = it:read()
end


print('Process Crosswikis')

it, _ = io.open(opt.root_data_dir .. 'basic_data/p_e_m_data/crosswikis_p_e_m.txt')
line = it:read()

while (line) do
  local parts = split(line, "\t")
  local mention  = parts[1]

  if (not mention:find('Wikipedia')) and (not mention:find('wikipedia')) then
    if not merged_e_m_counts[mention] then
      merged_e_m_counts[mention] = tds.Hash()
    end  

    local total_freq = tonumber(parts[2])
    assert(total_freq)
    local num_ents = table_len(parts)
    for i = 3, num_ents do
      local ent_str = split(parts[i], ",")
      local ent_wikiid = tonumber(ent_str[1])
      assert(ent_wikiid)
      local freq = tonumber(ent_str[2])
      assert(freq)

      if not merged_e_m_counts[mention][ent_wikiid] then
        merged_e_m_counts[mention][ent_wikiid] = 0
      end
      merged_e_m_counts[mention][ent_wikiid] = merged_e_m_counts[mention][ent_wikiid] + freq
    end
  end
  line = it:read()
end


print('Now sorting and writing ..')
out_file = opt.root_data_dir .. 'generated/crosswikis_wikipedia_p_e_m.txt'
ouf = assert(io.open(out_file, "w"))

for mention, list in pairs(merged_e_m_counts) do
  if mention:len() >= 1 then
    local tbl = {}
    for ent_wikiid, freq in pairs(list) do
      table.insert(tbl, {ent_wikiid = ent_wikiid, freq = freq})
    end
    table.sort(tbl, function(a,b) return a.freq > b.freq end)

    local str = ''
    local total_freq = 0
    local num_ents = 0
    for _,el in pairs(tbl) do
      if is_valid_ent(el.ent_wikiid) then
        str = str .. el.ent_wikiid .. ',' .. el.freq 
        str = str .. ',' .. get_ent_name_from_wikiid(el.ent_wikiid):gsub(' ', '_') .. '\t'
        num_ents = num_ents + 1
        total_freq = total_freq + el.freq

        if num_ents >= 100 then -- At most 100 candidates
          break
        end
      end
    end
    ouf:write(mention .. '\t' .. total_freq .. '\t' .. str .. '\n')
  end
end
ouf:flush()
io.close(ouf)

print('    Done sorting and writing.')



<|start_filename|>data_gen/gen_p_e_m/gen_p_e_m_from_wiki.lua<|end_filename|>
-- Generate p(e|m) index from Wikipedia
-- Run: th data_gen/gen_p_e_m/gen_p_e_m_from_wiki.lua -root_data_dir $DATA_PATH

cmd = torch.CmdLine()
cmd:option('-root_data_dir', '', 'Root path of the data, $DATA_PATH.')
cmd:text()
opt = cmd:parse(arg or {})
assert(opt.root_data_dir ~= '', 'Specify a valid root_data_dir path argument.')


require 'torch'
dofile 'utils/utils.lua'
dofile 'data_gen/parse_wiki_dump/parse_wiki_dump_tools.lua'
tds = tds or require 'tds' 

print('\nComputing Wikipedia p_e_m')

it, _ = io.open(opt.root_data_dir .. 'basic_data/textWithAnchorsFromAllWikipedia2014Feb.txt')
line = it:read()

wiki_e_m_counts = tds.Hash()

-- Find anchors, e.g. <a href="wikt:anarchism">anarchism</a>
local num_lines = 0
local parsing_errors = 0
local list_ent_errors = 0
local diez_ent_errors = 0
local disambiguation_ent_errors = 0
local num_valid_hyperlinks = 0

while (line) do
  num_lines = num_lines + 1
  if num_lines % 5000000 == 0 then
    print('Processed ' .. num_lines .. ' lines. Parsing errs = ' .. 
          parsing_errors .. ' List ent errs = ' .. 
          list_ent_errors .. ' diez errs = ' .. diez_ent_errors .. 
          ' disambig errs = ' .. disambiguation_ent_errors .. 
          ' . Num valid hyperlinks = ' .. num_valid_hyperlinks)
  end

  if not line:find('<doc id="') then
    list_hyp, text, le_errs, p_errs, dis_errs, diez_errs = extract_text_and_hyp(line, false)
    parsing_errors = parsing_errors + p_errs
    list_ent_errors = list_ent_errors + le_errs
    disambiguation_ent_errors = disambiguation_ent_errors + dis_errs
    diez_ent_errors = diez_ent_errors + diez_errs
    for _,el in pairs(list_hyp) do
      local mention = el.mention
      local ent_wikiid = el.ent_wikiid
      
      -- A valid (entity,mention) pair
      num_valid_hyperlinks = num_valid_hyperlinks + 1

      if not wiki_e_m_counts[mention] then
        wiki_e_m_counts[mention] = tds.Hash()
      end
      if not wiki_e_m_counts[mention][ent_wikiid] then
        wiki_e_m_counts[mention][ent_wikiid] = 0
      end
      wiki_e_m_counts[mention][ent_wikiid] = wiki_e_m_counts[mention][ent_wikiid] + 1
    end
  end
  
  line = it:read()
end

print('    Done computing Wikipedia p(e|m). Num valid hyperlinks = ' .. num_valid_hyperlinks)

print('Now sorting and writing ..')
out_file = opt.root_data_dir .. 'generated/wikipedia_p_e_m.txt'
ouf = assert(io.open(out_file, "w"))

for mention, list in pairs(wiki_e_m_counts) do
  local tbl = {}
  for ent_wikiid, freq in pairs(list) do
    table.insert(tbl, {ent_wikiid = ent_wikiid, freq = freq})
  end
  table.sort(tbl, function(a,b) return a.freq > b.freq end)

  local str = ''
  local total_freq = 0
  for _,el in pairs(tbl) do
    str = str .. el.ent_wikiid .. ',' .. el.freq  
    str = str .. ',' .. get_ent_name_from_wikiid(el.ent_wikiid):gsub(' ', '_') .. '\t'
    total_freq = total_freq + el.freq
  end
  ouf:write(mention .. '\t' .. total_freq .. '\t' .. str .. '\n')
end
ouf:flush()
io.close(ouf)

print('    Done sorting and writing.')
