<commit_before>! Copyright (C) 2012 <NAME>.
! See http://factorcode.org/license.txt for BSD license.
USING: arrays assocs fry kernel math math.functions
math.statistics memoize random sequences sorting ;
IN: machine-learning.rebalancing

ERROR: probability-sum-not-one seq ;

: check-probabilities ( seq -- seq )
    dup sum 1.0 .00000000001 ~ [ probability-sum-not-one ] unless ;

: equal-probabilities ( n -- array )
    dup recip <array> ; inline

MEMO: probabilities-seq ( seq -- seq' )
    check-probabilities [ >float ] map cum-sum ;

: probabilities-quot ( seq -- quot )
    probabilities-seq
    '[ _ random-unit '[ _ > ] find drop ] ; inline

: stratified-sample ( stratified-sequences probability-sequence -- elt )
    probabilities-quot call swap nth random ; inline

: equal-stratified-sample ( stratified-sequences -- elt )
    random random ; inline

: balance-labels ( X y n -- X' y' )
    [
        dup [ ] collect-index-by
        values '[
            _ _ _ equal-stratified-sample
            '[ _ swap nth ] bi@ 2array
        ]
    ] dip swap replicate [ keys ] [ values ] bi ;

: skew-labels ( X y probs n -- X' y' )
    [
        [ dup [ ] collect-index-by sort-keys values ] dip
        '[
            _ _ _ _ stratified-sample
            '[ _ swap nth ] bi@ 2array
        ]<commit_msg>machine-learning: Add stratified-samples for getting multiple samples.<commit_after>! Copyright (C) 2012 <NAME>.
! See http://factorcode.org/license.txt for BSD license.
USING: arrays assocs fry kernel math math.functions
math.statistics memoize random sequences sorting ;
IN: machine-learning.rebalancing

ERROR: probability-sum-not-one seq ;

: check-probabilities ( seq -- seq )
    dup sum 1.0 .00000000001 ~ [ probability-sum-not-one ] unless ;

: equal-probabilities ( n -- array )
    dup recip <array> ; inline

MEMO: probabilities-seq ( seq -- seq' )
    check-probabilities [ >float ] map cum-sum ;

: probabilities-quot ( seq -- quot )
    probabilities-seq
    '[ _ random-unit '[ _ > ] find drop ] ; inline

: stratified-sample ( stratified-sequences probability-sequence -- elt )
    probabilities-quot call swap nth random ; inline

: stratified-samples ( stratified-sequences probability-sequence n -- elt )
    [ '[ _ _ stratified-sample ] ] dip swap replicate ;

: equal-stratified-sample ( stratified-sequences -- elt )
    random random ; inline

: balance-labels ( X y n -- X' y' )
    [
        dup [ ] collect-index-by
        values '[
            _ _ _ equal-stratified-sample
            '[ _ swap nth ] bi@ 2array
        ]
    ] dip swap replicate [ keys ] [ values ] bi ;

: skew-labels ( X y probs n -- X' y' )
    [
        [ dup [ ] collect-index-by sort-keys values ] dip
        '[
            _ _ _ _ stratified-sample
            '[ _ swap nth ] bi@ 2array
        ]
{'fraction_non_alphanumeric': 0.11549395877754087, 'fraction_numerical': 0.013503909026297086, 'mean_word_length': 2.3551847437425506, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, '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': '3310009', 'n_tokens_mistral': 891, 'n_tokens_neox': 833, 'n_words': 239}