<commit_before>module ResultSet (
    build 
  , ResultSet(..)
  , Results
) where

import Control.Parallel.Strategies
import Control.Monad.ST
import Data.Maybe (isJust)
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Vector.Algorithms.Intro (sort)
import Scorer
import Utils (merge, chunkV)

type ScoredList = V.Vector (Double, B.ByteString)

data Results = ParVec [ScoredList] 
               deriving (Show, Eq)

class ResultSet a where
  size   :: a -> Int
  items  :: a -> [B.ByteString]
  refine :: ScoreStrategy s => a -> s -> a

instance ResultSet Results where
  size  (ParVec sl) = sum . fmap V.length $ sl
  items (ParVec sl) = fmap snd . merge fst . fmap V.toList $ sl

  refine (ParVec sl) sc = ParVec newSet
    where rl     = fmap (fmap snd) sl
          newSet = scoreRL (score sc) rl

-- Create
build :: [B.ByteString] -> Results
build lst = ParVec chunks
  where initVec   = V.fromList $ zip [1..] lst
        len       = length lst
        chunkSize = fst . divMod len $ 5000
        chunks    = chunkV (chunkSize + 1) $ initVec
  
-- Score line accordingly
scoreRL :: Scorer -> [V.Vector B.ByteString] -> [ScoredList]
scoreRL f rl = parMap rdeepseq cms rl
  where fo x = fmap (\i -> (i, x)) $ f x
        cms x = runST $ do
              let remaining = V.filter isJust . fmap fo $ x
              let vsize = V.length remaining
              -- Copy the array to a mutable one
              mv <- MV.new vsize

              V.forM_ (V.indexed remaining) $ \res -> do
                  case res of
                    (idx, Just el) -> MV.write mv idx el
                    _              -> return ()
                  
              -- Sort
              sort mv
              V.unsafeFreeze mv
<commit_msg>indexM<commit_after>module ResultSet (
    build 
  , ResultSet(..)
  , Results
) where

import Control.Parallel.Strategies
import Control.Monad
import Control.Monad.ST
import Data.Maybe (isJust)
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Vector.Algorithms.Intro (sort)
import Scorer
import Utils (merge, chunkV)

type ScoredList = V.Vector (Double, B.ByteString)

data Results = ParVec [ScoredList] 
               deriving (Show, Eq)

class ResultSet a where
  size   :: a -> Int
  items  :: a -> [B.ByteString]
  refine :: ScoreStrategy s => a -> s -> a

instance ResultSet Results where
  size  (ParVec sl) = sum . fmap V.length $ sl
  items (ParVec sl) = fmap snd . merge fst . fmap V.toList $ sl

  refine (ParVec sl) sc = ParVec newSet
    where rl     = fmap (fmap snd) sl
          newSet = scoreRL (score sc) rl

-- Create
build :: [B.ByteString] -> Results
build lst = ParVec chunks
  where initVec   = V.fromList $ zip [1..] lst
        len       = length lst
        chunkSize = fst . divMod len $ 5000
        chunks    = chunkV (chunkSize + 1) $ initVec
  
-- Score line accordingly
scoreRL :: Scorer -> [V.Vector B.ByteString] -> [ScoredList]
scoreRL f rl = parMap rdeepseq cms rl
  where fo x = fmap (\i -> (i, x)) $ f x
        cms x = runST $ do
              let remaining = V.filter isJust . fmap fo $ x
              let vsize = V.length remaining
              -- Copy the array to a mutable one
              mv <- MV.new vsize

              forM_ [0..(vsize - 1)] $ \idx -> do
                  e <- V.indexM remaining idx
                  case e of
                    Just el -> MV.write mv idx el
                    _       -> return ()
                  
              -- Sort
              sort mv
              V.unsafeFreeze mv

{'fraction_non_alphanumeric': 0.09321340964840556, 'fraction_numerical': 0.004360861270100845, 'mean_word_length': 1.8209069946195235, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7414142', 'n_tokens_mistral': 1167, 'n_tokens_neox': 1129, 'n_words': 405}