swift-lda 0.4.1 → 0.7.0.0
raw patch · 2 files changed
+97/−39 lines, 2 files
Files
- NLP/SwiftLDA.hs +96/−38
- swift-lda.cabal +1/−1
NLP/SwiftLDA.hs view
@@ -9,10 +9,42 @@ -- topics), even though it is generic. For example if used for word -- class induction, replace documents with word types, words with -- features and topics with word classes.+--+-- Usage example:+-- +-- > import qualified Data.Vector as V+-- > import qualified Data.Vector.Unboxed as U+-- > import Control.Monad.ST+-- > import NLP.SwiftLDA+-- > +-- > main = do+-- > -- Initialize model.+-- > let river = 1+-- > money = 2+-- > bank = 3+-- > m <- stToIO $ initial (U.singleton 42) 2 10 0.1+-- > let docs = [ (1, U.fromList [river, river, bank]) +-- > , (2, U.fromList [money, money, bank])+-- > , (3, U.fromList [river, money, bank])+-- > ]+-- > -- Run 10 iterations of sampling on this batch of documents.+-- > result <- stToIO $ run m 10 docs+-- > -- Display topic assignments.+-- > print result+-- > -- Run one iteration of sampling on another batch and display result+-- > print =<< (stToIO $ run m 1 [(4, U.fromList [bank, bank])]) +-- > -- Retrieve and display word-topic weights.+-- > fm <- stToIO $ finalize m+-- > print $ wordTopics fm+-- > + module NLP.SwiftLDA- ( -- * Samplers- pass- , passOne+ ( + -- * Initialization and finalization+ initial+ , finalize+ -- * Running sampler+ , run -- * Datatypes , LDA , Doc@@ -23,9 +55,6 @@ , Table1D -- * Access model information , Finalized (..) - -- * Initialization and finalization- , initial- , finalize -- * Querying evolving model , docTopicWeights_ , priorDocTopicWeights_@@ -50,6 +79,7 @@ import qualified Data.IntMap as IntMap import qualified Data.List as List import qualified Data.Foldable as Fold+import qualified Data.Traversable as Trav import GHC.Generics (Generic) import Debug.Trace @@ -62,11 +92,17 @@ type Table2D = IntMap.IntMap Table1D type Table1D = IntMap.IntMap Double +-- | Document ID type D = Int+-- | Topic ID type Z = Int +-- | Word ID type W = Int-type Doc = (D, U.Vector (W, Maybe Z))+-- | A document consists of a document ID and a sequence of word IDs.+type Doc = (D, U.Vector W) +type SampledDoc = (D, U.Vector (W, Maybe Z))+ -- | Abstract type holding the settings and the state of the sampler data LDA s = LDA @@ -103,13 +139,13 @@ , exponent :: !(Maybe Double) -- ^ Learning rate exponent } deriving (Generic) --- | For each document sum the topic counts +-- | For each document sum the topic counts. docCounts :: Finalized -> Table1D docCounts = IntMap.map (sum . IntMap.elems) . docTopics -- | Create transparent immutable object holding model information--- from opaque internal representation+-- from opaque internal representation. finalize :: LDA s -> ST s Finalized finalize m = do dt <- read . _docTopics $ m@@ -139,11 +175,14 @@ iWSIZE :: Int iWSIZE = 1000 --- | @initial s k a b@ initializes model with @k@ topics, @a/k@ alpha--- hyperparameter, @b@ beta hyperparameter and random seed @s@-initial :: U.Vector Word32 -> Int -> Double -> Double -> Maybe Double +-- | @initial s k a b@ initializes model with @k@ topics, @a/k@+-- alpha hyperparameter, @b@ beta hyperparameter and random seed @s@.+initial :: U.Vector Word32 + -> Int + -> Double + -> Double -> ST s (LDA s) -initial s k a b e = do +initial s k a b = do dta <- newArray ((0,0),(iDSIZE, k-1)) 0 wta <- newArray ((0,0),(iWSIZE, k-1)) 0 LDA <$> @@ -157,18 +196,40 @@ newArray (0,k-1) 0 <*> new 0 <*> initialize s <*>- pure e+ pure Nothing rho :: Double -> Int -> Double rho e t = 1 - (1 + fromIntegral t)**(-e) {-# INLINE rho #-} --- | @pass batch@ runs one pass of Gibbs sampling on documents in @batch@ -pass :: Int -> LDA s -> V.Vector Doc -> ST s (V.Vector Doc)-pass t m = V.mapM (passOne t m) +-- | @run m i batch@ runs an outer loop of @i@ passes of Gibbs+-- sampling over documents in @batch@ using the model @m@ and returns+-- the topic assignments for words in the documents of the batch.+run :: (Trav.Traversable f) => + LDA s+ -> Int + -> f Doc+ -> ST s (f (U.Vector Z))+run m i docs = do + let prepareDoc (!d, ws) = (d, U.map (\ !w -> (w, Nothing)) ws)+ {-# INLINE prepareDoc #-}+ assignments = U.map (maybe (error "NLP.SwiftLDA.run: Nothing") id . snd) + {-# INLINE assignments #-}+ go !z !j = pass 1 m z+ {-# INLINE go #-}+ sampled <- Fold.foldlM go (fmap prepareDoc docs) $ [1..i]+ return $! fmap (assignments . snd) sampled+{-# SPECIALIZE run :: LDA s -> Int -> V.Vector Doc -> ST s (V.Vector (U.Vector Z)) #-}+ +-- | @pass t m batch@ runs one pass of Gibbs sampling on documents in+-- @batch@ with @t@ as time index.+pass :: (Trav.Traversable f) => Int -> LDA s -> f SampledDoc -> ST s (f SampledDoc)+pass t m = Trav.mapM (passOne t m) +{-# SPECIALIZE pass :: Int -> LDA s -> V.Vector SampledDoc -> ST s (V.Vector SampledDoc) #-} --- | Run a pass on a single doc -passOne :: Int -> LDA s -> Doc -> ST s Doc++-- | Run a pass on a single doc. +passOne :: Int -> LDA s -> SampledDoc -> ST s SampledDoc passOne t m doc@(!d, wz) = do grow m doc zs <- U.mapM one wz@@ -182,7 +243,7 @@ update r m d w z return z --- | Sample a random topic for doc d and word w +-- | Sample a random topic for doc d and word w. randomZ :: LDA s -> Int -> Int -> ST s Int randomZ m !d !w = do wordTopicWeights_ m d w@@ -192,7 +253,7 @@ -- | @topicWeights m d w@ sets @weights m@ to the unnormalized probabilities of -- topics for word @w@ in document @d@ given LDA model @m@. --- Each call overwrites current weights+-- Each call overwrites current weights. wordTopicWeights_ :: LDA s -> Int -> Int -> ST s () wordTopicWeights_ m !d !w = do let k = _topicNum m@@ -216,27 +277,26 @@ go (z+1) go l +priorDocTopicWeights_ :: LDA s -> D -> ST s (U.Vector Double)+priorDocTopicWeights_ m d = do+ grow m (d, U.empty)+ dt <- read (_docTopics m)+ ((_,0),(_,u)) <- getBounds dt+ U.generateM (u+1) (\z -> readArray dt (d,z)) + -- Weight calc on Finalized -- | @docTopicWeights m doc@ returns unnormalized topic probabilities--- for document doc given LDA model @m@+-- for document doc given LDA model @m@. docTopicWeights :: Finalized -> Doc -> U.Vector Double docTopicWeights m (d, ws) = U.accumulate (+) (U.replicate (topicNum m) 0) . U.concatMap (U.indexed . wordTopicWeights m d)- . U.map fst $ ws {-# INLINE docTopicWeights #-}- -priorDocTopicWeights_ :: LDA s -> D -> ST s (U.Vector Double)-priorDocTopicWeights_ m d = do- grow m (d, U.empty)- dt <- read (_docTopics m)- ((_,0),(_,u)) <- getBounds dt- U.generateM (u+1) (\z -> readArray dt (d,z)) docTopicWeights_ :: LDA s -> Doc -> ST s (U.Vector Double) docTopicWeights_ m doc@(d, ws) = do- grow m doc+ grow m (d, U.map (\w -> (w, Nothing)) ws) (0,u) <- getBounds (weights m) let r = U.replicate (_topicNum m) 0 let one w = do@@ -244,7 +304,6 @@ U.generateM (u+1) (readArray (weights m)) U.foldM' (\z w -> do y <- one w return $! U.zipWith (+) z y) r - . U.map fst $ ws -- | @topicWeights m d w@ returns the unnormalized probabilities of -- topics for word @w@ in document @d@ given LDA model @m@.@@ -259,7 +318,7 @@ weights = [ (count z dt + a) -- n(z,d) + alpha * (count z wt + b) -- n(z,w) + beta * (1/(count z (topics m) + v * b)) - -- n(.,w) + V * beta+ -- n(z,.) + V * beta | z <- [0..k-1] ] in U.fromList weights {-# INLINE wordTopicWeights #-}@@ -270,13 +329,13 @@ dt <- read (_docTopics m) wt <- read (_wordTopics m) wsz <- read (_wSize m) ; write (_wSize m) (max (w+1) wsz)- nz <- readArray (_topics m) z ; writeArray (_topics m) z (nz+c)- nzd <- readArray dt (d,z) ; writeArray dt (d,z) (nzd+c)- nzw <- readArray wt (w,z) ; writeArray wt (w,z) (nzw+c)+ nz <- readArray (_topics m) z ; writeArray (_topics m) z (max 0 $ nz+c)+ nzd <- readArray dt (d,z) ; writeArray dt (d,z) (max 0 $ nzd+c)+ nzw <- readArray wt (w,z) ; writeArray wt (w,z) (max 0 $ nzw+c) -- | @grow m doc@ grows the @docTopic m@ and @wordTopic m@ tables as necessary -- according to content of @doc@-grow :: LDA s -> Doc -> ST s ()+grow :: LDA s -> SampledDoc -> ST s () grow m (d, wz) = do let w = if U.null wz then 0 else U.maximum (U.map fst wz) dt <- read (_docTopics m) ; (_,(d_max,_)) <- getBounds dt @@ -319,7 +378,6 @@ go (i+1) (n-v) | otherwise = return (i-1) go l r- read :: STRef s a -> ST s a read = readSTRef
swift-lda.cabal view
@@ -1,5 +1,5 @@ Name: swift-lda-Version: 0.4.1+Version: 0.7.0.0 Synopsis: Online sampler for Latent Dirichlet Allocation Description: Online Gibbs sampler for Latent Dirichlet Allocation. LDA is a generative admixture model frequently used