def link_documents(self,
                     max_n=None,
                     fact_out_file="/tmp/facts.json",
                     qry_out_file="/tmp/queries.json",
                     para_out_file="/tmp/paragraphs.json",
                     filter_subjects=None,
                     exclude_subjects=None):
    """Load n documents and link them to facts."""
    start = time.time()
    fout = open(fact_out_file, "w")
    fq_out, fp_out = open(qry_out_file, "w"), open(para_out_file, "w")
    seen_articles = set()
    total_paras = 0
    for n, (doc_id, doc_raw) in enumerate(self.corpus.input):
      doc_id = str(doc_id, "utf-8")
      if n % 1000 == 0:
        print("processed", n, "items in %.1f" % (time.time() - start), "sec")
      if max_n is not None and random.uniform(0, 1) > (float(max_n) / 550000):
        continue
      if filter_subjects is not None and doc_id not in filter_subjects:
        continue
      if exclude_subjects is not None and doc_id in exclude_subjects:
        continue
      # get kb items
      seen_articles.add(doc_id)
      kb_item = self.kb[doc_id]
      tail_entities = {}
      all_properties = []
      for prop, tail in kb_item:
        tup = self.get_canonical_property(prop, tail)
        if tup is None:
          tup = self.get_date_property(prop, tail)
          if tup is None:
            continue
        tail_entities[tup[1]] = tup[0]
        all_properties.append(tup[0])
      store = sling.Store(self.commons)
      document = sling.Document(store.parse(doc_raw), store, self.docschema)
      if not document.tokens:
        print("Skipping %s No tokens." % (doc_id))
        continue
      # build token maps
      tok_to_sent_id, tok_to_para_id, sent_to_span, para_to_span = {}, {}, {}, {}
      tok_to_char_offset = {}
      offset = 0
      sent_begin = para_begin = 0
      sent_id = para_id = 0
      for ii, token in enumerate(document.tokens):
        if ii > 0 and token.brk == 4:
          para_to_span[para_id] = (para_begin, ii)
          sent_to_span[sent_id] = (sent_begin, ii)
          para_id += 1
          sent_id += 1
          sent_begin = para_begin = ii
        elif ii > 0 and token.brk == 3:
          sent_to_span[sent_id] = (sent_begin, ii)
          sent_id += 1
          sent_begin = ii
        tok_to_sent_id[ii] = sent_id
        tok_to_para_id[ii] = para_id
        tok_to_char_offset[ii] = offset
        offset += len(token.word) + 1
      para_to_span[para_id] = (para_begin, len(document.tokens))
      sent_to_span[sent_id] = (sent_begin, len(document.tokens))
      # find subjects
      sent_to_subj, para_to_subj = (collections.defaultdict(list),
                                    collections.defaultdict(list))
      para_to_ment = collections.defaultdict(list)
      ment_to_para_index = {}
      mentid_to_linked_entity = {}
      sorted_mentions = sorted(document.mentions, key=lambda m: m.begin)
      for ii, mention in enumerate(sorted_mentions):
        if tok_to_sent_id[mention.begin] != tok_to_sent_id[mention.end - 1]:
          continue
        linked_entity = self.get_linked_entity(mention)
        mentid_to_linked_entity[ii] = linked_entity
        para_id = tok_to_para_id[mention.begin]
        para_to_ment[para_id].append((mention, linked_entity))
        ment_to_para_index[ii] = len(para_to_ment[para_id]) - 1
        if linked_entity == doc_id:
          sent_id = tok_to_sent_id[mention.begin]
          sent_to_subj[sent_id].append(ii)
          para_to_subj[para_id].append(ii)

      # save paragraphs
      local_to_global_para = {}
      for para_id, para_span in para_to_span.items():
        if para_span[1] - para_span[0] < MIN_LEN:
          continue
        if len(para_to_ment[para_id]) <= 1:
          continue
        local_to_global_para[para_id] = total_paras
        fp_out.write(
            self.serialize_para(document, tok_to_char_offset, para_span,
                                para_to_ment[para_id], doc_id, total_paras) +
            "\n")
        total_paras += 1

      # find tails
      seen_properties = {}
      for ii, mention in enumerate(sorted_mentions):
        # first look for sentence matches
        linked_entity = mentid_to_linked_entity[ii]
        if linked_entity == doc_id:
          continue
        if linked_entity in tail_entities:
          if tail_entities[linked_entity] in seen_properties:
            continue
          my_sent = tok_to_sent_id[mention.begin]
          my_para = tok_to_para_id[mention.begin]
          para_span = para_to_span[my_para]
          if my_para not in local_to_global_para:
            continue
          if my_sent in sent_to_subj:
            # sent_span = sent_to_span[my_sent]
            fq_out.write(
                self.serialize_query(
                    local_to_global_para[my_para], ment_to_para_index, doc_id,
                    sent_to_subj[my_sent], tail_entities[linked_entity],
                    linked_entity, ii, "sentence") + "\n")
            subj_mentions = [
                para_to_ment[my_para][ment_to_para_index[mm]][0]
                for mm in sent_to_subj[my_sent]
            ]
            fout.write(
                self.serialize_fact(document, tok_to_char_offset, para_span,
                                    doc_id, subj_mentions,
                                    tail_entities[linked_entity], linked_entity,
                                    mention, "sentence") + "\n")
            seen_properties[tail_entities[linked_entity]] = my_para
            self.relation_stats["sentences"][tail_entities[linked_entity]] += 1

      for ii, mention in enumerate(sorted_mentions):
        # next look for paragraph matches
        linked_entity = mentid_to_linked_entity[ii]
        if linked_entity == doc_id:
          continue
        if linked_entity in tail_entities:
          if tail_entities[linked_entity] in seen_properties:
            continue
          my_para = tok_to_para_id[mention.begin]
          para_span = para_to_span[my_para]
          if my_para not in local_to_global_para:
            continue
          if my_para in para_to_subj:
            fq_out.write(
                self.serialize_query(
                    local_to_global_para[my_para], ment_to_para_index, doc_id,
                    para_to_subj[my_para], tail_entities[linked_entity],
                    linked_entity, ii, "paragraph") + "\n")
            subj_mentions = [
                para_to_ment[my_para][ment_to_para_index[mm]][0]
                for mm in para_to_subj[my_para]
            ]
            fout.write(
                self.serialize_fact(document, tok_to_char_offset, para_span,
                                    doc_id, subj_mentions,
                                    tail_entities[linked_entity], linked_entity,
                                    mention, "paragraph") + "\n")
            seen_properties[tail_entities[linked_entity]] = my_para
            self.relation_stats["paragraphs"][tail_entities[linked_entity]] += 1

      # add negatives
      max_neg = len(seen_properties)
      num_neg = 0
      all_para_id = list(para_to_subj.keys())
      if not all_para_id:
        continue
      for tail, prop in tail_entities.items():
        if num_neg == max_neg:
          break
        if prop in seen_properties:
          continue
        random_para_id = random.choice(all_para_id)
        random_para_span = para_to_span[random_para_id]
        subj_mentions = [
            para_to_ment[random_para_id][ment_to_para_index[mm]][0]
            for mm in para_to_subj[random_para_id]
        ]
        fout.write(
            self.serialize_fact(document, tok_to_char_offset, random_para_span,
                                doc_id, subj_mentions, prop, None, None,
                                "entity negative") + "\n")
        num_neg += 1
        seen_properties[prop] = None
        self.relation_stats["entity negatives"][prop] += 1

    fout.close()
    fq_out.close()
    fp_out.close()
    print("Sentences -- ", sum(self.relation_stats["sentences"].values()))
    print(" :: ".join(
        "%s:%d" % (k, v) for k, v in self.relation_stats["sentences"].items()))
    print("Paragraphs -- ", sum(self.relation_stats["paragraphs"].values()))
    print(" :: ".join(
        "%s:%d" % (k, v) for k, v in self.relation_stats["paragraphs"].items()))
Load n documents and link them to facts.
language/labs/drkit/wikidata/preprocessing/distantly_supervise.py