<commit_before># Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
"""A script to build the tf-idf document matrices for retrieval.

Adapted from <NAME>'s work at github.com/facebookresearch/DrQA/
"""

import torch
import numpy as np
import scipy.sparse as sp
import argparse
import os
import math
import logging

from multiprocessing import Pool as ProcessPool
from multiprocessing.util import Finalize
from functools import partial
from collections import Counter

from . import utils
from .doc_db import DocDB
from drqa import tokenizers

logger = logging.getLogger()
logger.setLevel(logging.INFO)
fmt = logging.Formatter('%(asctime)s: [ %(message)s ]', '%m/%d/%Y %I:%M:%S %p')
console = logging.StreamHandler()
console.setFormatter(fmt)
logger.addHandler(console)


# ------------------------------------------------------------------------------
# Multiprocessing functions
# ------------------------------------------------------------------------------

PROCESS_TOK = None
PROCESS_DB = None


def init(tokenizer_class, db_opts):
    global PROCESS_TOK, PROCESS_DB
    PROCESS_TOK = tokenizer_class()
    Finalize(PROCESS_TOK, PROCESS_TOK.shutdown, exitpriority=100)
    PROCESS_DB = DocDB(**db_opts)
    Finalize(PROCESS_DB, PROCESS_DB.close, exitpriority=100)


def fetch_text(doc_id):
    global PROCESS_DB
    return PROCESS_DB.get_doc_text(doc_id)


def tokenize(text):
    global PROCESS_TOK
    return PROCESS_TOK.tokenize(text)


MAX_SZ = math.pow(2, 31)

# ------------------------------------------------------------------------------
# Build article --> word count sparse matrix.
# ------------------------------------------------------------------------------


def truncate(data, row, col):
    global MAX_SZ
    if len(data) > MAX_SZ:
        over = len(data) - MAX_SZ
        pct = over / len(data)
        logger.info('Data size is too large for scipy to index all of it. '
                    'Throwing out {} entries ({}%% of data).'.format(over, pct))
        data = data[:MAX_SZ]
        row = row[:MAX_SZ]
        col = col[:MAX_SZ]
    return data, row, col

def sparse_nonzero(sparse_t):
    return sparse_t.coalesce()._indices()
<commit_msg>fix to retriever max size (#524)<commit_after># Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
"""A script to build the tf-idf document matrices for retrieval.

Adapted from <NAME>'s work at github.com/facebookresearch/DrQA/
"""

import torch
import numpy as np
import scipy.sparse as sp
import argparse
import os
import math
import logging

from multiprocessing import Pool as ProcessPool
from multiprocessing.util import Finalize
from functools import partial
from collections import Counter

from . import utils
from .doc_db import DocDB
from drqa import tokenizers

logger = logging.getLogger()
logger.setLevel(logging.INFO)
fmt = logging.Formatter('%(asctime)s: [ %(message)s ]', '%m/%d/%Y %I:%M:%S %p')
console = logging.StreamHandler()
console.setFormatter(fmt)
logger.addHandler(console)


# ------------------------------------------------------------------------------
# Multiprocessing functions
# ------------------------------------------------------------------------------

PROCESS_TOK = None
PROCESS_DB = None


def init(tokenizer_class, db_opts):
    global PROCESS_TOK, PROCESS_DB
    PROCESS_TOK = tokenizer_class()
    Finalize(PROCESS_TOK, PROCESS_TOK.shutdown, exitpriority=100)
    PROCESS_DB = DocDB(**db_opts)
    Finalize(PROCESS_DB, PROCESS_DB.close, exitpriority=100)


def fetch_text(doc_id):
    global PROCESS_DB
    return PROCESS_DB.get_doc_text(doc_id)


def tokenize(text):
    global PROCESS_TOK
    return PROCESS_TOK.tokenize(text)


MAX_SZ = int(math.pow(2, 30) * 1.8)

# ------------------------------------------------------------------------------
# Build article --> word count sparse matrix.
# ------------------------------------------------------------------------------


def truncate(data, row, col):
    global MAX_SZ
    if len(data) > MAX_SZ:
        over = len(data) - MAX_SZ
        pct = over / len(data)
        logger.info('Data size is too large for scipy to index all of it. '
                    'Throwing out {} entries ({}%% of data).'.format(over, pct))
        data = data[:MAX_SZ]
        row = row[:MAX_SZ]
        col = col[:MAX_SZ]
    return data, row, col

def sparse_nonzero(sparse_t):
    return sparse_t.coalesce()._indices()

{'fraction_non_alphanumeric': 0.21199011124845488, 'fraction_numerical': 0.00638648537288834, 'mean_word_length': 2.238825883922615, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1680906', 'n_tokens_mistral': 1449, 'n_tokens_neox': 1371, 'n_words': 434}