sequor-0.2: src/Perceptron/Sequence.hs
{-# LANGUAGE NoMonomorphismRestriction
, BangPatterns
, FlexibleInstances
#-}
module Perceptron.Sequence
(
Model(..)
, Options(..)
, Eval
, YMap
, train
, decode
)
where
import Data.Array.ST
import Data.Array.Unboxed
import qualified Data.Array as A
import qualified Data.Vector.Unboxed as V
import Control.Monad.ST
import Data.STRef
import Control.Monad
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import Perceptron.Vector
import System.IO
import Debug.Trace
import Config
import Data.List (inits,foldl',sortBy)
import Data.Ord (comparing)
import ListZipper
import qualified Data.Binary as Binary
import Utils (uniq)
data Model = Model { options :: Options
, weights :: UArray I Float }
type X = [Xi]
type Y = [Yi]
type I = (Yi,Xii)
type Xi = V.Vector Xii
type Xii = Int
type Yi = Int
type Dot = LocalSparseVector Yi Xii -> Float
data Options = Options { oYMap :: YMap
, oIndexSet :: IntSet.IntSet
, oYDict :: IntMap.IntMap [Yi]
, oYs :: [Yi]
, oBeam :: Int
, oRate :: Float
, oEpochs :: Int
, oFeatBounds :: Maybe (Int,Int)
} deriving Eq
type YMap = (Xi,A.Array Yi Xi,A.Array (Yi,Yi) Xi)
instance Binary.Binary (V.Vector Int) where
put v = Binary.put $ V.toList v
get = V.fromList `fmap` Binary.get
instance Binary.Binary Model where
put m = do
Binary.put (options m)
-- Binary.put (weights m)
let (lo,hi) = bounds . weights $ m
xs = filter (\(_,e) -> e /= 0.0) . assocs . weights $ m
Binary.put (lo,hi)
Binary.put xs
get = {-# SCC "get1" #-} do
os <- Binary.get
os == os `seq` return ()
ws <- do
(lo,hi) <- Binary.get
xs <- Binary.get
xs == xs `seq` return ()
return $ accumArray (+) 0 (lo,hi) $ xs
ws == ws `seq` return ()
return $ Model os ws
instance Binary.Binary Options where
put (Options a b c d e f g h) = Binary.put a >> Binary.put b >> Binary.put c
>> Binary.put d >> Binary.put e >> Binary.put f
>> Binary.put g >> Binary.put h
get = {-# SCC "get2" #-} do
a <- Binary.get
a == a `seq` return ()
b <- Binary.get
b == b `seq` return ()
c <- Binary.get
c == c `seq` return ()
d <- Binary.get
d == d `seq` return ()
e <- Binary.get
e == e `seq` return ()
f <- Binary.get
f == f `seq` return ()
g <- Binary.get
g == g `seq` return ()
h <- Binary.get
h == h `seq` return ()
return $ Options a b c d e f g h
yDictFind :: Options -> Xi -> [Yi]
yDictFind opts fs =
let mk = V.find (`IntSet.member` oIndexSet opts) $ fs
def = oYs opts
in case mk of
Just k -> IntMap.findWithDefault def k . oYDict $ opts
Nothing -> def
-- | DECODING
decode :: Model -> X -> Y
decode m = fst . decode' (options m) (weights m `dot`)
data Cell = Cell { cScore :: !Float
, cPhi :: SparseVector I
, cPath :: Y
, cStep :: ListZipper Xi } deriving (Show,Eq)
decode' :: Options -> Dot -> X -> (Y,SparseVector I)
decode' opts w x =
bestPath opts w [Cell { cScore = 0
, cPhi = Map.empty
, cPath = []
, cStep = fromList x } ]
phi :: Options -> X -> Y -> SparseVector I
phi opts x y = foldl' f Map.empty . zip x . map reverse . tail . inits $ y
where f z (xi,ys) = z `plus` toSV (features (oYMap opts) xi ys)
{-# INLINE features #-}
features :: YMap -> Xi -> [Yi] -> LocalSparseVector Yi Xii
features (!zero,uni,bi) xi (y:ys) =
case ys of
[] -> (y, zero V.++ xi)
[y1] -> (y, uni A.! y1 V.++ xi)
(y1 : y2 : _) -> let r = bi A.! (y1,y2)
in if V.null r
then (y, uni A.! y1 V.++ xi)
else (y, r V.++ xi)
beamSearch :: Options
-> Dot
-> [Cell]
-> [Cell]
beamSearch opts w cs =
let f cs = if any (atEnd . cStep) cs then cs
else
let cs' = [ let fs = features (oYMap opts) xi (y':ys)
in Cell { cScore =
s + w fs
, cPhi = ph `plus` (toSV fs)
, cPath = (y':ys)
, cStep = next x }
| Cell { cScore = s
, cPhi = ph
, cPath = ys
, cStep = x } <- cs
, let Just xi = focus x::Maybe Xi
, y' <- yDictFind opts xi
]
in f . take (oBeam opts)
. sortBy (flip $ comparing cScore)
$ cs'
in f cs
bestPath :: Options
-> Dot
-> [Cell]
-> (Y, SparseVector I)
bestPath opts w xs =
let xs' = beamSearch opts w xs
first = (\(x:_) -> x) xs'
in ( reverse . cPath $ first
, cPhi first )
-- | TRAINING
iter :: Options
-> Int
-> [(X,Y)]
-> (STRef s Int, DenseVectorST s I, DenseVectorST s I)
-> ST s ()
iter opts _ ss (c,params,params_a) = do
for_ ss $ \ (x,y) -> do
params' <- unsafeFreeze params
let (y',phi_xy') = decode' opts (params'`dot`) x
when (y' /= y) $ do
let phi_xy = phi opts x y
update = (phi_xy `minus` phi_xy') `scale` oRate opts
params `plus_` update
c' <- readSTRef c
params_a `plus_` (update `scale` fromIntegral c')
modifySTRef c (+1)
type Eval = Int -> [(Y,Y)] -> [(Y,Y)] -> String
train :: Options -> [(X, Y)] -> Eval -> [(X,Y)] -> Model
train opts heldout eval ss = Model opts $ runSTUArray $ do
let bs = computeBounds opts ss
trace ("Param vector bounds: " ++ show bs) () `seq` return ()
params <- newArray bs 0
params_a <- newArray bs 0
c <- newSTRef 1
let undef = error "Perceptron.Sequence.train: undefined"
runLogger . hPutStrLn stderr $ eval 0 undef undef
for_ [1..oEpochs opts] $
\i -> do iter opts i ss (c,params,params_a)
c' <- readSTRef c
params' <- unsafeFreeze params
params_a' <- unsafeFreeze params_a
let w = (fromIntegral c',params',params_a')
ys xys = [ fst . decode' opts (w`dot'`) $ x
| (x,_) <- xys ]
runLogger
. hPutStrLn stderr
$ eval i (zip (map snd ss) (ys ss))
(zip (map snd heldout) (ys heldout))
finalParams (c, params, params_a)
return params
{-# NOINLINE runLogger #-}
runLogger f = unsafeIOToST f
finalParams :: (STRef s Int, DenseVectorST s I, DenseVectorST s I)
-> ST s ()
finalParams (c,params,params_a) = do
(l,u) <- getBounds params
c' <- fmap fromIntegral (readSTRef c)
for_ (range (l,u)) $ \i -> do
e <- readArray params i
e_a <- readArray params_a i
writeArray params i (e - (e_a * (1/c')))
computeBounds :: Options -> [(X,Y)] -> (I,I)
computeBounds opts xys =
let ((yl,xl),(yh,xh)) = foldl' f ((maxBound,minimum xis)
,(minBound,maximum xis))
. (\(xs,ys) -> zip (concat xs) (concat ys))
. unzip
$ xys
in case oFeatBounds opts of
Just (xl',xh') -> ((yl,xl'),(yh,xh'))
Nothing -> ((yl,xl),(yh,xh))
where f ((!miny,!minx),(!maxy,!maxx)) (xs,!y) =
((min miny y,V.minimum $ minx`V.cons`xs)
,(max maxy y,V.maximum $ maxx`V.cons`xs))
xis = let (zero,uni,bi) = oYMap opts
in uniq
. concatMap V.toList
$
[zero]
++
(filter (not . V.null)
. A.elems
$ bi)
++
(filter (not . V.null)
. A.elems
$ uni)