<commit_before>package cc.factorie.app.nlp.mention

import scala.collection.mutable.HashSet
import scala.collection.mutable.ArrayBuffer
import cc.factorie.app.nlp._
import cc.factorie.app.nlp.parse.ParseTree
import cc.factorie.app.nlp.ner.NerSpan
import cc.factorie.app.nlp.pos.PTBPosLabel
import scala.Some

object ParseBasedMentionFinding extends DocumentAnnotator {

  private final val PERSONAL_PRONOUNS = Seq("PRP", "PRP$")
  private final val COMMON_NOUNS      = Seq("NN" , "NNS")
  private final val PROPER_NOUNS      = Seq("NNP", "NNPS")
  private final val ALL_NOUNS         = Seq("NN","NNS","NNP","NNPS","PRP","PRP$")

  private def isPersonalPronoun(t: Token) = PERSONAL_PRONOUNS exists(_ == t.posLabel.categoryValue.toUpperCase())
  private def isCommonNoun     (t: Token) = COMMON_NOUNS      exists(_ == t.posLabel.categoryValue.toUpperCase())
  private def isProperNoun     (t: Token) = PROPER_NOUNS      exists(_ == t.posLabel.categoryValue.toUpperCase())
  private def isNoun           (t: Token) = ALL_NOUNS         exists(_ == t.posLabel.categoryValue.toUpperCase())

  private def nerSpans(doc: Document): Seq[Mention] = {
    (for (s <- doc.spansOfClass[NerSpan]) yield
      Mention(doc, s.start, s.length,s.start - s.sentence.start  + s.length - 1) //this sets the head token idx to be the last token in the span  //todo: s.start is offset in document, right?
      ).toSeq
  }

  // [Assumes personal pronouns are single tokens.]
  private def personalPronounSpans(doc: Document): Seq[Mention] = {
    (
      for (s <- doc.sentences;
           (t,i) <- s.tokens.zipWithIndex if isPersonalPronoun(t)) yield
        Mention(doc, s.start + i, 1,s.start+i)
      ).toSeq
  }

  private def getHeadTokenIdx(m: Mention): Int = {
    getHead(
      m.document.sentenceContaining(m.document(m.start)).parse,
      m.start until (m.start + m.length) // TODO: is the until correct here?
    )
  }

  def getHead(parse: ParseTree, subtree: Seq[Int]): Int = {
    var curr = subtree.head - parse.sentence.start
    while (parse.parentIndex(curr) > 0 && subtree.contains(parse.parentIndex(curr)))
      curr = parse.parentIndex(curr)
    curr + parse.sentence.start
  }<commit_msg>bug fix and efficiency improvements to parse-based mention finder.<commit_after>package cc.factorie.app.nlp.mention

import scala.collection.mutable.HashSet
import scala.collection.mutable.ArrayBuffer
import cc.factorie.app.nlp._
import cc.factorie.app.nlp.parse.ParseTree
import cc.factorie.app.nlp.ner.NerSpan
import cc.factorie.app.nlp.pos.PTBPosLabel
import scala.Some

object ParseBasedMentionFinding extends DocumentAnnotator {

  private final val PERSONAL_PRONOUNS = Seq("PRP", "PRP$")
  private final val COMMON_NOUNS      = Seq("NN" , "NNS")
  private final val PROPER_NOUNS      = Seq("NNP", "NNPS")
  private final val ALL_NOUNS         = Seq("NN","NNS","NNP","NNPS","PRP","PRP$")

  private def isPersonalPronoun(t: Token) = PERSONAL_PRONOUNS exists(_ == t.posLabel.categoryValue.toUpperCase())
  private def isCommonNoun     (t: Token) = COMMON_NOUNS      exists(_ == t.posLabel.categoryValue.toUpperCase())
  private def isProperNoun     (t: Token) = PROPER_NOUNS      exists(_ == t.posLabel.categoryValue.toUpperCase())
  private def isNoun           (t: Token) = ALL_NOUNS         exists(_ == t.posLabel.categoryValue.toUpperCase())

  private def nerSpans(doc: Document): Seq[Mention] = {
    (for (s <- doc.spansOfClass[NerSpan]) yield
      Mention(doc, s.start, s.length,s.start - s.sentence.start  + s.length - 1) //this sets the head token idx to be the last token in the span  //todo: s.start is offset in document, right?
      ).toSeq
  }

  // [Assumes personal pronouns are single tokens.]
  private def personalPronounSpans(doc: Document): Seq[Mention] = {
    (
      for (s <- doc.sentences;
           (t,i) <- s.tokens.zipWithIndex if isPersonalPronoun(t)) yield
        Mention(doc, s.start + i, 1,s.start+i)
      ).toSeq
  }

  private def getHeadTokenIdx(m: Mention): Int = {
    getHead(
      m.document.sentenceContaining(m.document(m.start)).parse,
      m.start until (m.start + m.length) // TODO: is the until correct here?
    )
  }
  //this expects as input indices in the **document** not the sentence
  //note that this never returns the root as the head, it always returns a pointer to an actual token in the sentence
  //it will either return the root of a parse tree span, or a token that is a child of the root
  def getHead(parse: ParseTree, subtree: Seq[Int]): Int = {
    val sentenceLevelIndices = subtree.map(i => i - parse.sentence.start)
    var curr = sentenceLevelIndices.head
    val leftBoundary = sentenceLevelIndices.head
    val rightBoundary = sentenceLevelIndices.last
    while(parse.parentIndex(curr) > 0 && containedInInterval(leftBoundary,rightBoundary,parse.parentIndex(curr))){
      curr = parse.parentIndex(curr)
    }
    curr + parse.sentence.start  //this shifts it back to have document-level indices
  }

  private def containedInInterval(left: Int, right: Int, testIndex: Int): Boolean = {
    testIndex >= left && testIndex <= right
  }
{'fraction_non_alphanumeric': 0.11876345137937781, 'fraction_numerical': 0.0011739385638818234, 'mean_word_length': 3.5399644760213143, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7059753', 'n_tokens_mistral': 1727, 'n_tokens_neox': 1688, 'n_words': 432}