packages feed

colada-0.0.1: Colada/WordClass.hs

{-# LANGUAGE 
   OverloadedStrings  
 , FlexibleInstances
 , DeriveGeneric 
 , NoMonomorphismRestriction 
 , DeriveDataTypeable
 , TemplateHaskell
 #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | Word Class induction with LDA
--
-- This module provides function which implement word class induction
-- using the generic algorithm implemented in Colada.LDA. 
--
-- You can access and set options in the @Options@ record using lenses.
-- Example:
--
-- >  import Data.Label
-- >  let options =   set passes 5 
-- >                . set beta 0.01 
-- >                . set topicNum 100 
-- >                $ defaultOptions
-- >  in run options sentences

module Colada.WordClass 
       ( 
         -- * Running 
         run
       , defaultOptions
       , summary
         -- * Class and word prediction
       , label
       , predict
         -- * Data types and associated lenses
       , WordClass       
       , ldaModel
         -- | LDA model
       , atomTable
         -- | String to atom and vice versa conversion tables
       , options
         -- | Options for Gibbs sampling
       , Options
       , featIds
         -- | Feature ids
       , topicNum
         -- | Number of topics K
       , alphasum
         -- | Dirichlet parameter alpha*K which controls topic sparseness
       , beta
         -- | Dirichlet parameter beta which controls word sparseness
       , passes
         -- | Number of sampling passes per batch
       , repeats
         -- | Number of repeats per sentences
       , batchSize
         -- | Number of sentences per batch
       , seed
         -- | Random seed for the sampler
       )
where       
  
-- Standard libraries  
import qualified Data.Text.Lazy as Text
import qualified Data.Text.Lazy.Encoding as Text
import qualified Data.Vector  as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed  as U
import qualified Data.IntMap as IntMap
import qualified Data.Serialize as Serialize
import qualified Control.Monad as M
import qualified Data.List as List
import qualified Data.List.Split as Split
import qualified Data.Ord as Ord
import Data.Word (Word64)
import Data.Typeable (Typeable)
import Data.Data (Data)
import Prelude hiding ((.))
import Control.Category ((.))

-- Third party modules  
import qualified Control.Monad.Atom  as Atom
import qualified NLP.CoNLL  as CoNLL
import qualified Data.List.Zipper as Zipper
import GHC.Generics (Generic)
import qualified Data.Label as L
import Data.Label (get)

import qualified NLP.LDA as LDA
import NLP.LDA.Utils (count)

-- Package modules
import qualified Colada.Features as F


-- | Container for the Word Class model
data WordClass = 
  WordClass { _ldaModel   :: LDA.Finalized      -- ^ LDA model
            , _atomTable  :: Atom.AtomTable (U.Vector Char) -- ^ String to
                                                      -- Int
                                                      -- conversion
                                                      -- table
            , _options    :: Options 
            }
  deriving (Generic)

data Options = Options { _featIds  :: [Int]
                       , _topicNum :: !Int     
                       , _alphasum :: !Double  
                       , _beta     :: !Double  
                       , _passes   :: !Int     
                       , _repeats  :: !Int     
                       , _batchSize :: !Int    
                       , _seed      :: !Word64 
                       }
             deriving (Eq, Show, Typeable, Data, Generic)

instance Serialize.Serialize Options

$(L.mkLabels [''WordClass, ''Options])

defaultOptions :: Options
defaultOptions = Options { _featIds = [-1,1]
                         , _topicNum  = 10                
                         , _alphasum  = 10
                         , _beta      = 0.1
                         , _passes    = 1
                         , _repeats   = 1
                         , _batchSize = 1 
                         , _seed      = 0 
                         }
                 
-- | @run options xs@ runs the LDA Gibbs sampler for word classes with
-- @options@ on sentences @xs@, and returns the resulting model
run :: Options -> [CoNLL.Sentence] -> WordClass
run opts xs = 
  let (ss, atomTab) = flip Atom.runAtom Atom.empty 
                          . prepareData (get repeats opts)
                                        (get featIds opts)
                          $ xs
      bs = batches (get batchSize opts) ss                    
      m  = LDA.initial (get topicNum opts) (get alphasum opts) (get beta opts)
      lda = mapM_ (\b -> M.foldM (const . LDA.pass) b [1..get passes opts]) 
                  bs
      m' = snd . LDA.runSampler (get seed opts) m $ lda
  in WordClass (LDA.finalize m') atomTab opts 



-- | @prepareData rep is ss@ replicates each sentence in stream @ss@
-- @rep@ times. Features with indices @is@ are extracted from each
-- token, and word and features are converted to ints in the Atom
-- Monad.
prepareData ::  Int 
             -> [Int] 
             -> [CoNLL.Sentence]
             -> Atom.Atom (U.Vector Char) [V.Vector LDA.Doc]
prepareData rep is ss = do
  let mk fs = let d = IntMap.findWithDefault 
                      (error "parseData: focus feature missing") 0 fs
                  ws =   [ Text.concat [f,"^",Text.pack . show $ i ] 
                         | i <- is , Just f <- [IntMap.lookup i fs] ]
               in (d, ws)
      doc (d, ws) = do
        da <- Atom.toAtom . compress $ d
        was <- mapM (Atom.toAtom . compress) ws
        return (da, U.fromList $ zip was (repeat Nothing))
      sent s = do fs <- mapM (doc . mk) . extractFeatures  $ s
                  return $! V.fromList fs
  ss' <- mapM sent ss
  return $! concatMap (replicate rep) ss'
     
-- | @batches sz ss@ creates batches of size @sz@ from the stream of
-- sentence feature vectors @ss@. The vectors in a batch are
-- concatenated.
batches :: Int -> [V.Vector LDA.Doc] -> [V.Vector LDA.Doc]
batches sz = map V.concat . Split.chunk sz

-- | @summary m@ returns a textual summary of word classes found in
-- model @m@
summary :: WordClass -> Text.Text
summary m = 
  let format (z,cs) =  do 
        cs' <-  mapM (Atom.fromAtom . fst)
                . takeWhile ((>0) . snd)
                . take 10 
                . List.sortBy (flip $ Ord.comparing snd) . IntMap.toList 
                $ cs
        return . Text.unwords $ Text.pack (show z) 
          : map (Text.pack . U.toList) cs' 
  in fst . flip Atom.runAtom (get atomTable m) 
    . M.liftM Text.unlines 
    . mapM format
    . IntMap.toList
    . IntMap.fromListWith (IntMap.unionWith (+))
    . concatMap (\(d,zs) -> [ (z, IntMap.singleton d c) | (z,c) <- zs ])
    . IntMap.toList
    . IntMap.map  IntMap.toList
    . LDA.docTopics
    . LDA.model
    . get ldaModel 
    $ m


-- | @label m s@ returns for each word in sentences s, unnormalized
-- probabilities of word classes.
label :: WordClass -> CoNLL.Sentence -> V.Vector (U.Vector Double)
label m s = fst . Atom.runAtom label' . L.get atomTable $ m
  where label' = do
          let fm = L.get ldaModel m
          s' <- prepareSent  m s
          return $! V.map (LDA.docTopicWeights . LDA.model $ fm) $ s'
              
-- | @predict m s@ returns for each word in sentence s, unnormalized
-- probabilities of words given predicted word class.
predict :: WordClass -> CoNLL.Sentence 
           -> V.Vector (V.Vector (Double, Text.Text))
predict m s = fst . Atom.runAtom predict' . L.get atomTable $ m
  where predict' = do
          let fm = L.get ldaModel m
          s' <- prepareSent m s
          let ws = V.map  (G.convert . predictDoc fm . docToWs) 
                   $ s'
          V.mapM (V.mapM fromAtom) ws
        docToWs = U.map fst . snd
        fromAtom (n,w) = do w' <- Atom.fromAtom w
                            return (n, decompress w')

prepareSent :: WordClass -> CoNLL.Sentence 
               -> Atom.Atom (U.Vector Char) (V.Vector LDA.Doc)
prepareSent m s = do 
  [r] <- prepareData 1 (L.get (featIds . options) m) [s]
  return r

-- | @predictDoc m ws@ returns unnormalized probabilities of each
-- document id given the model @m@ and words @ws@. The candidate
-- document ids are taken from the model @m@.  The weights are
-- computed according to the following formula:
--  
-- > P(d|{w}) ∝ Σ_z[n(d,z)+a Σ_{w in ws}(n(w,z)+b)/(Σ_{w in V} n(w,z)+b)]
predictDoc ::  LDA.Finalized -> U.Vector LDA.W -> U.Vector (Double, LDA.D)
predictDoc m ws =
  let k = LDA.topicNum . LDA.model $ m
      a = LDA.alphasum (LDA.model m) / fromIntegral k
      b = LDA.beta . LDA.model $ m
      v = fromIntegral . LDA.vSize . LDA.model $ m
      zt = LDA.topics . LDA.model $ m
      wsums = 
        IntMap.fromList 
          [ (z, U.sum . U.map (\w -> (count w wt_z + b) / denom) $ ws)
          | z <- IntMap.keys zt
          , let wt_z = IntMap.findWithDefault IntMap.empty z . LDA.topicWords 
                       $ m
                denom  = count z zt + b * v ]
      wsum z = IntMap.findWithDefault 
                (error "Colada.LDA.predictDoc: key not found")
                z wsums
  in   U.fromList 
     . List.sortBy (flip $ Ord.comparing id) 
     $ [ ( sum [ (c + a) * wsum z | (z,c) <- IntMap.toList zt_d ] , d)
       | d <- IntMap.keys . LDA.docTopics . LDA.model $ m
       , let zt_d = IntMap.findWithDefault IntMap.empty d 
                    . LDA.docTopics 
                    . LDA.model
                    $ m ]

extractFeatures :: CoNLL.Sentence -> [IntMap.IntMap Text.Text]
extractFeatures = F.featureSeq combine  
                  . Zipper.fromList 
                  . map (V.! 0) 
                  . V.toList
  
combine :: Maybe Text.Text -> Maybe Text.Text -> Maybe Text.Text
combine (Just a) (Just b) = Just $ Text.concat [a, "|", b]
combine (Just a) Nothing  = Just a
combine Nothing  (Just b) = Just b
combine Nothing Nothing   = Nothing
  



compress :: Text.Text -> U.Vector Char
compress = U.fromList . Text.unpack

decompress :: U.Vector Char -> Text.Text
decompress = Text.pack . U.toList

-- Instances for serialization

instance Serialize.Serialize LDA.Finalized
instance Serialize.Serialize LDA.LDA

instance Serialize.Serialize Text.Text where
  put = Serialize.put . Text.encodeUtf8
  get = Text.decodeUtf8 `fmap` Serialize.get
  
instance Serialize.Serialize (U.Vector Char) where
  put v = Serialize.put (Text.pack . U.toList $ v)
  get = 
    do t <- Serialize.get
       return $! U.fromList . Text.unpack $ t
       
instance Serialize.Serialize (Atom.AtomTable (U.Vector Char))

instance Serialize.Serialize WordClass