<commit_before>package wikidump

import (
	"bytes"
	"golang.org/x/text/unicode/norm"
	"html"
	"regexp"
	"strings"
	"unicode"
	"unicode/utf8"
)

var (
	starttag = regexp.MustCompile(`^<[a-z][^>]*>`)
	endtag   = regexp.MustCompile(`^</[a-z]+>`)
)

// Get rid of tables, template calls, quasi-XML. Throws away their content.
//
// Assumes tables, templates and tags are properly nested, except for spurious
// end-of-{table,template,element} tags, which are ignored.
func Cleanup(s string) string {
	var depth int
	output := bytes.NewBuffer(make([]byte, 0, len(s)))

	for {
		next := strings.IndexAny(s, "{}|<")
		if next == -1 {
			if depth == 0 {
				output.WriteString(s)
			}
			break
		}

		if depth == 0 {
			output.WriteString(s[:next])
		}
		s = s[next:]

		var skip int
		if strings.HasPrefix(s, "{{") || strings.HasPrefix(s, "{|") {
			depth++
			skip = 2
		} else if strings.HasPrefix(s, "|}") || strings.HasPrefix(s, "}}") {
			if depth > 0 {
				depth--
			}
			skip = 2
		} else if s[0] != '<' {
			// This case prevents regexp matching for a 20% speedup.
			skip = 1
		} else if span := starttag.FindStringIndex(s); span != nil {
			depth++
			skip = span[1]
		} else if span := endtag.FindStringIndex(s); span != nil {
			if depth > 0 {
				depth--
			}
			skip = span[1]
		} else {
			skip = 1
		}

		// If skip == 1, we didn't find a tag/table marker.
		if skip == 1 && depth == 0 && len(s) > 0 {
			output.WriteByte(s[0])
		}
		s = s[skip:]
	}
	return norm.NFC.String(html.UnescapeString(output.String()))
}

// A link to the article Target with anchor text Anchor.
type Link struct {
	Anchor, Target string
}

var (
	linkRE     = regexp.MustCompile(`(\w*)\[\[([^]]+)\]\](\w*)`)
	whitespace = regexp.MustCompile(`[\s_]+`)
)

func normSpace(s string) string {
	s = whitespace.ReplaceAllLiteralString(s, " ")
	return strings.TrimSpace(s)
}

// Extract all the wikilinks from s. Returns a frequency table.
func ExtractLinks(s string) map[Link]int {
	freq := make(map[Link]int)

	for _, candidate := range linkRE.FindAllStringSubmatch(s, -1) {
		before, l, after := candidate[1], candidate[2], candidate[3]

		var target, anchor string
		if pipe := strings.IndexByte(l, '|'); pipe != -1 {
			target, anchor = l[:pipe], l[pipe+1:]
		} else {
			target = l
			anchor = l
		}

		// If the anchor contains a colon, assume it's a file or category link.
		// XXX Maybe skip matches for `:\s`? Proper solution would parse the
		// dump to find non-main namespace prefixes.
		if strings.IndexByte(target, ':') != -1 {
			continue
		}
<commit_msg>Speed up link extraction by some 20%.<commit_after>package wikidump

import (
	"bytes"
	"golang.org/x/text/unicode/norm"
	"html"
	"regexp"
	"strings"
	"unicode"
	"unicode/utf8"
)

var (
	starttag = regexp.MustCompile(`^<[a-z][^>]*>`)
	endtag   = regexp.MustCompile(`^</[a-z]+>`)
)

// Get rid of tables, template calls, quasi-XML. Throws away their content.
//
// Assumes tables, templates and tags are properly nested, except for spurious
// end-of-{table,template,element} tags, which are ignored.
func Cleanup(s string) string {
	var depth int
	output := bytes.NewBuffer(make([]byte, 0, len(s)))

	for {
		next := strings.IndexAny(s, "{}|<")
		if next == -1 {
			if depth == 0 {
				output.WriteString(s)
			}
			break
		}

		if depth == 0 {
			output.WriteString(s[:next])
		}
		s = s[next:]

		var skip int
		if strings.HasPrefix(s, "{{") || strings.HasPrefix(s, "{|") {
			depth++
			skip = 2
		} else if strings.HasPrefix(s, "|}") || strings.HasPrefix(s, "}}") {
			if depth > 0 {
				depth--
			}
			skip = 2
		} else if s[0] != '<' {
			// This case prevents regexp matching for a 20% speedup.
			skip = 1
		} else if span := starttag.FindStringIndex(s); span != nil {
			depth++
			skip = span[1]
		} else if span := endtag.FindStringIndex(s); span != nil {
			if depth > 0 {
				depth--
			}
			skip = span[1]
		} else {
			skip = 1
		}

		// If skip == 1, we didn't find a tag/table marker.
		if skip == 1 && depth == 0 && len(s) > 0 {
			output.WriteByte(s[0])
		}
		s = s[skip:]
	}
	return norm.NFC.String(html.UnescapeString(output.String()))
}

// A link to the article Target with anchor text Anchor.
type Link struct {
	Anchor, Target string
}

func normSpace(s string) string {
	parts := strings.FieldsFunc(s, func(c rune) bool {
		return c == '_' || unicode.IsSpace(c)
	})
	return strings.Join(parts, " ")
}

var linkRE = regexp.MustCompile(`\w*\[\[[^]]+\]\]\w*`)

// Extract all the wikilinks from s. Returns a frequency table.
func ExtractLinks(s string) map[Link]int {
	freq := make(map[Link]int)

	for _, candidate := range linkRE.FindAllStringSubmatch(s, -1) {
		// Parse complex links like "foo[[bar|baz]]quux". We used to do this
		// with capturing groups in linkRE, but those are *slow*.
		text := candidate[0]
		openbrack := strings.IndexByte(text, '[')
		closebrack := strings.LastIndex(text, "]")
		before := text[:openbrack]
		after := text[closebrack+1:]
		mid := text[openbrack+2:closebrack-1]

		var target, anchor string
		if pipe := strings.IndexByte(mid, '|'); pipe != -1 {
			target, anchor = mid[:pipe], mid[pipe+1:]
		} else {
			target = mid
			anchor = mid
		}

		// If the anchor contains a colon, assume it's a file or category link.
		// XXX Maybe skip matches for `:\s`? Proper solution would parse the
		// dump to find non-main namespace prefixes.
		if strings.IndexByte(target, ':') != -1 {
			continue
		}

{'fraction_non_alphanumeric': 0.1677144970414201, 'fraction_numerical': 0.010909763313609468, 'mean_word_length': 3.0946252838758515, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 15, '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': '734581', 'n_tokens_mistral': 2201, 'n_tokens_neox': 1928, 'n_words': 547}