1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
|
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
module ProbabilityMap
( fileName,
ProbabilityCache (..),
readProbabilityCache,
writeProbabilityCache,
modifyProbabilityCache,
fromDocuments,
addDocument,
deleteDocument,
replaceDocument,
)
where
import Control.Applicative (liftA2)
import Control.Arrow ((***))
import Control.DeepSeq (NFData (..))
import Control.Exception (SomeException, try)
import Control.Parallel.Strategies (parList, rdeepseq, withStrategy)
import Data.Binary qualified as B
import Data.List (foldl')
import Data.Map qualified as M
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Data.Text qualified as T
import Data.Vector qualified as V
import Data.Vector.Binary ()
import Document qualified as D
import GHC.Conc (numCapabilities)
import GHC.Generics (Generic)
import Parallel (streamsOf)
import Tag qualified as G
data ProbabilityCache = ProbabilityCache
{ docsPerWord :: M.Map T.Text Int,
docsPerTag :: M.Map G.Tag Int,
docsPerWordAndTag :: M.Map (T.Text, G.Tag) Int,
probabilityMap :: M.Map (T.Text, G.Tag) Double
}
deriving (Show, Generic)
data ProbabilityCache' = ProbabilityCache'
{ docsPerWord :: M.Map T.Text Int,
docsPerTag :: M.Map G.Tag Int,
docsPerWordAndTag :: M.Map (T.Text, G.Tag) Int
}
deriving (Show, Generic)
instance NFData ProbabilityCache' where
rnf (ProbabilityCache' {..}) =
rnf docsPerWord `seq`
rnf docsPerTag `seq`
rnf docsPerWordAndTag
fromProbabilityCache' :: ProbabilityCache' -> ProbabilityCache
fromProbabilityCache' probabilityCache'@(ProbabilityCache' {..}) =
ProbabilityCache {probabilityMap = probabilityMap probabilityCache', ..}
toProbabilityCache' :: ProbabilityCache -> ProbabilityCache'
toProbabilityCache' (ProbabilityCache {..}) =
ProbabilityCache' {..}
instance B.Binary ProbabilityCache where
get = fromCacheRep <$> B.get
put = B.put . toCacheRep
instance NFData ProbabilityCache where
rnf (ProbabilityCache {..}) =
rnf docsPerWord `seq`
rnf docsPerTag `seq`
rnf docsPerWordAndTag `seq`
rnf probabilityMap
data CacheRep = CacheRep
{ allWords :: V.Vector T.Text,
allTags :: V.Vector G.Tag,
docsPerWord :: M.Map Int Int,
docsPerTag :: M.Map Int Int,
docsPerWordAndTag :: M.Map (Int, Int) Int,
probabilityMap :: M.Map (Int, Int) Double
}
deriving (Show, Generic, B.Binary)
instance NFData CacheRep where
rnf (CacheRep {..}) =
rnf allWords `seq`
rnf allTags `seq`
rnf docsPerWord `seq`
rnf docsPerTag `seq`
rnf docsPerWordAndTag `seq`
rnf probabilityMap
fromCacheRep :: CacheRep -> ProbabilityCache
fromCacheRep (CacheRep {..}) =
ProbabilityCache
{ docsPerWord = M.mapKeys fromWord docsPerWord,
docsPerTag = M.mapKeys fromTag docsPerTag,
docsPerWordAndTag = M.mapKeys fromWordAndTag docsPerWordAndTag,
probabilityMap = M.mapKeys fromWordAndTag probabilityMap
}
where
fromWord = (allWords V.!)
fromTag = (allTags V.!)
fromWordAndTag = fromWord *** fromTag
toCacheRep :: ProbabilityCache -> CacheRep
toCacheRep (ProbabilityCache {..}) =
CacheRep
{ allWords = V.fromList allWords',
allTags = V.fromList allTags',
docsPerWord = M.mapKeys toWord docsPerWord,
docsPerTag = M.mapKeys toTag docsPerTag,
docsPerWordAndTag = M.mapKeys toWordAndTag docsPerWordAndTag,
probabilityMap = M.mapKeys toWordAndTag probabilityMap
}
where
allWords' = M.keys docsPerWord
allTags' = M.keys docsPerTag
allWordsMap = M.fromList $ zip allWords' [0 ..]
allTagsMap = M.fromList $ zip allTags' [0 ..]
toWord = (allWordsMap M.!)
toTag = (allTagsMap M.!)
toWordAndTag = toWord *** toTag
fileName :: FilePath
fileName = "probabilityCache"
readProbabilityCache :: IO ProbabilityCache
readProbabilityCache =
try (B.decodeFile fileName) >>= \case
Left (_ :: SomeException) -> fromDocuments <$> D.getDocuments
Right probabilityCache -> pure probabilityCache
writeProbabilityCache :: ProbabilityCache -> IO ()
writeProbabilityCache probabilityCache =
B.encodeFile fileName probabilityCache
modifyProbabilityCache :: (ProbabilityCache -> ProbabilityCache) -> IO ProbabilityCache
modifyProbabilityCache f =
liftA2 (>>) (writeProbabilityCache . f) pure =<< readProbabilityCache
fromDocuments :: [D.Document] -> ProbabilityCache
fromDocuments allDocs =
fromProbabilityCache' $
foldl' (flip addDocument') (ProbabilityCache' M.empty M.empty M.empty) allDocs
probabilityMap :: ProbabilityCache' -> M.Map (T.Text, G.Tag) Double
probabilityMap (ProbabilityCache' {..}) =
M.unions
. withStrategy (parList rdeepseq)
. map
( M.fromList
. filter (\(_, p) -> p > 0)
. map (\k@(word, tag) -> (k, probability word tag))
)
$ streamsOf
numCapabilities
[ (word, tag)
| word <- allWords,
tag <- allTags
]
where
probability word tag =
let docsWithWord = docsPerWord M.! word
docsWithTag = docsPerTag M.! tag
docsWithWordAndTag = fromMaybe 0 $ M.lookup (word, tag) docsPerWordAndTag
in fi docsWithWordAndTag / fi (docsWithWord + docsWithTag - docsWithWordAndTag)
allWords = M.keys docsPerWord
allTags = M.keys docsPerTag
fi = fromIntegral @Int @Double
addDocument :: D.Document -> ProbabilityCache -> ProbabilityCache
addDocument doc probabilityCache =
fromProbabilityCache' $ addDocument' doc (toProbabilityCache' probabilityCache)
deleteDocument :: D.Document -> ProbabilityCache -> ProbabilityCache
deleteDocument doc probabilityCache =
fromProbabilityCache' $ deleteDocument' doc (toProbabilityCache' probabilityCache)
replaceDocument :: D.Document -> D.Document -> ProbabilityCache -> ProbabilityCache
replaceDocument oldDoc newDoc probabilityCache =
fromProbabilityCache' $
(addDocument' newDoc . deleteDocument' oldDoc) (toProbabilityCache' probabilityCache)
addDocument' :: D.Document -> ProbabilityCache' -> ProbabilityCache'
addDocument' doc probabilityCache =
combining (+) probabilityCache (fromDocument doc)
deleteDocument' :: D.Document -> ProbabilityCache' -> ProbabilityCache'
deleteDocument' doc probabilityCache =
combining (-) probabilityCache (fromDocument doc)
combining :: (Int -> Int -> Int) -> ProbabilityCache' -> ProbabilityCache' -> ProbabilityCache'
combining f probabilityCache probabilityCache' =
ProbabilityCache' {..}
where
docsPerWord = M.unionWith f probabilityCache.docsPerWord probabilityCache'.docsPerWord
docsPerTag = M.unionWith f probabilityCache.docsPerTag probabilityCache'.docsPerTag
docsPerWordAndTag = M.unionWith f probabilityCache.docsPerWordAndTag probabilityCache'.docsPerWordAndTag
fromDocument :: D.Document -> ProbabilityCache'
fromDocument doc =
let words = S.toList doc.index.originalWords
tags = S.toList doc.index.tags
wordAndTags = [(word, tag) | word <- words, tag <- tags]
in ProbabilityCache'
{ docsPerWord = docsPer words,
docsPerTag = docsPer tags,
docsPerWordAndTag = docsPer wordAndTags
}
where
docsPer :: Ord k => [k] -> M.Map k Int
docsPer = M.fromList . map (flip (,) 1)
|
