{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TupleSections #-} -- Inference with CRFs. module Data.CRF.Chain1.Constrained.Inference ( tag , tagK , marginals , accuracy , expectedFeaturesIn , zx , zx' ) where import Control.Applicative ((<$>)) import Data.Maybe (catMaybes) import Data.List (maximumBy, sortBy) import Data.Function (on) import qualified Data.Array as A import qualified Data.Vector as V import qualified Data.Vector.Unboxed as U import Control.Parallel.Strategies (rseq, parMap) import Control.Parallel (par, pseq) import GHC.Conc (numCapabilities) import qualified Data.Number.LogFloat as L import qualified Data.CRF.Chain1.Constrained.DP as DP import Data.CRF.Chain1.Constrained.Util (partition) import Data.CRF.Chain1.Constrained.Dataset.Internal import Data.CRF.Chain1.Constrained.Feature (Feature(..)) import Data.CRF.Chain1.Constrained.Model import Data.CRF.Chain1.Constrained.Intersect type LbIx = Int type ProbArray = Int -> LbIx -> L.LogFloat -- Some basic definitions. -- | Vector of potential labels on the given position of the sentence. lbVec :: Model -> Xs -> Int -> AVec Lb lbVec crf xs k = case xs V.! k of X _ -> (r0 crf) R _ r -> r {-# INLINE lbVec #-} -- | Number of potential labels on the given position of the sentence. lbNum :: Model -> Xs -> Int -> Int lbNum crf xs = (U.length . unAVec) . lbVec crf xs {-# INLINE lbNum #-} -- | Potential label on the given vector position. lbOn :: Model -> X -> Int -> Lb lbOn crf (X _) = (unAVec (r0 crf) U.!) lbOn _ (R _ r) = (unAVec r U.!) {-# INLINE lbOn #-} lbIxs :: Model -> Xs -> Int -> [(Int, Lb)] lbIxs crf xs = zip [0..] . U.toList . unAVec . lbVec crf xs {-# INLINE lbIxs #-} -- | Compute the table of potential products associated with -- observation features for the given sentence position. computePsi :: Model -> Xs -> Int -> LbIx -> L.LogFloat computePsi crf xs i = (A.!) $ A.accumArray (*) 1 bounds [ (k, valueL crf ix) | ob <- unX (xs V.! i) , (k, ix) <- intersect (obIxs crf ob) (lbVec crf xs i) ] where bounds = (0, lbNum crf xs i - 1) -- | Forward table computation. forward :: Model -> Xs -> ProbArray forward crf xs = alpha where alpha = DP.flexible2 (0, V.length xs) bounds (\t i -> withMem (computePsi crf xs i) t i) bounds i | i == V.length xs = (0, 0) | otherwise = (0, lbNum crf xs i - 1) withMem psi alpha i | i == V.length xs = const u | i == 0 = \j -> let x = lbOn crf (xs V.! i) j in psi j * sgValue crf x | otherwise = \j -> let x = lbOn crf (xs V.! i) j in psi j * ((u - v x) + w x) where u = sum [ alpha (i-1) k | (k, _) <- lbIxs crf xs (i-1) ] v x = sum [ alpha (i-1) k | (k, _) <- intersect (prevIxs crf x) (lbVec crf xs (i-1)) ] w x = sum [ alpha (i-1) k * valueL crf ix | (k, ix) <- intersect (prevIxs crf x) (lbVec crf xs (i-1)) ] -- | Backward table computation. backward :: Model -> Xs -> ProbArray backward crf xs = beta where beta = DP.flexible2 (0, V.length xs) bounds (\t i -> withMem (computePsi crf xs i) t i) bounds i | i == 0 = (0, 0) | otherwise = (0, lbNum crf xs (i-1) - 1) withMem psi beta i | i == V.length xs = const 1 | i == 0 = const $ sum [ beta (i+1) k * psi k * sgValue crf (lbOn crf (xs V.! i) k) | (k, _) <- lbIxs crf xs i ] | otherwise = \j -> let y = lbOn crf (xs V.! (i-1)) j in (u - v y) + w y where u = sum [ beta (i+1) k * psi k | (k, _ ) <- lbIxs crf xs i ] v y = sum [ beta (i+1) k * psi k | (k, _ ) <- intersect (nextIxs crf y) (lbVec crf xs i) ] w y = sum [ beta (i+1) k * psi k * valueL crf ix | (k, ix) <- intersect (nextIxs crf y) (lbVec crf xs i) ] zxBeta :: ProbArray -> L.LogFloat zxBeta beta = beta 0 0 zxAlpha :: Xs -> ProbArray -> L.LogFloat zxAlpha xs alpha = alpha (V.length xs) 0 -- | Normalization factor computed for the 'Xs' sentence using the -- backward computation. zx :: Model -> Xs -> L.LogFloat zx crf = zxBeta . backward crf -- | Normalization factor computed for the 'Xs' sentence using the -- forward computation. zx' :: Model -> Xs -> L.LogFloat zx' crf sent = zxAlpha sent (forward crf sent) -- | Tag probabilities with respect to marginal distributions. marginals :: Model -> Xs -> [[(Lb, L.LogFloat)]] marginals crf xs = let alpha = forward crf xs beta = backward crf xs in [ [ (x, prob1 alpha beta i k) | (k, x) <- lbIxs crf xs i ] | i <- [0 .. V.length xs - 1] ] -- | Get (at most) k best tags for each word and return them in -- descending order. TODO: Tagging with respect to marginal -- distributions might not be the best idea. Think of some -- more elegant method. tagK :: Int -> Model -> Xs -> [[(Lb, L.LogFloat)]] tagK k crf xs = map ( take k . reverse . sortBy (compare `on` snd) ) (marginals crf xs) -- | Find the most probable label sequence (with probabilities of individual -- lables determined with respect to marginal distributions) satisfying the -- constraints imposed over label values. tag :: Model -> Xs -> [Lb] tag crf = map (fst . head) . (tagK 1 crf) prob1 :: ProbArray -> ProbArray -> Int -> LbIx -> L.LogFloat prob1 alpha beta k x = alpha k x * beta (k + 1) x / zxBeta beta {-# INLINE prob1 #-} prob2 :: Model -> ProbArray -> ProbArray -> Int -> (LbIx -> L.LogFloat) -> LbIx -> LbIx -> FeatIx -> L.LogFloat prob2 crf alpha beta k psi x y ix = alpha (k - 1) y * beta (k + 1) x * psi x * valueL crf ix / zxBeta beta {-# INLINE prob2 #-} goodAndBad :: Model -> Xs -> Ys -> (Int, Int) goodAndBad crf xs ys = foldl gather (0, 0) $ zip labels labels' where labels = [ (best . unY) (ys V.! i) | i <- [0 .. V.length ys - 1] ] best zs | null zs = Nothing | otherwise = Just . fst $ maximumBy (compare `on` snd) zs labels' = map Just $ tag crf xs gather (good, bad) (x, y) | x == y = (good + 1, bad) | otherwise = (good, bad + 1) goodAndBad' :: Model -> [(Xs, Ys)] -> (Int, Int) goodAndBad' crf dataset = let add (g, b) (g', b') = (g + g', b + b') in foldl add (0, 0) [goodAndBad crf x y | (x, y) <- dataset] -- | Compute the accuracy of the model with respect to the labeled dataset. accuracy :: Model -> [(Xs, Ys)] -> Double accuracy crf dataset = let k = numCapabilities parts = partition k dataset xs = parMap rseq (goodAndBad' crf) parts (good, bad) = foldl add (0, 0) xs add (g, b) (g', b') = (g + g', b + b') in fromIntegral good / fromIntegral (good + bad) expectedFeaturesOn :: Model -> ProbArray -> ProbArray -> Xs -> Int -> [(FeatIx, L.LogFloat)] expectedFeaturesOn crf alpha beta xs i = tFeats ++ oFeats where psi = computePsi crf xs i pr1 = prob1 alpha beta i pr2 = prob2 crf alpha beta i psi oFeats = [ (ix, pr1 k) | o <- unX (xs V.! i) , (k, ix) <- intersect (obIxs crf o) (lbVec crf xs i) ] tFeats | i == 0 = catMaybes [ (, pr1 k) <$> featToIx crf (SFeature x) | (k, x) <- lbIxs crf xs i ] | otherwise = [ (ix, pr2 k l ix) | (k, x) <- lbIxs crf xs i , (l, ix) <- intersect (prevIxs crf x) (lbVec crf xs (i-1)) ] -- | A list of features (represented by feature indices) defined within -- the context of the sentence accompanied by expected probabilities -- determined on the basis of the model. -- -- One feature can occur multiple times in the output list. expectedFeaturesIn :: Model -> Xs -> [(FeatIx, L.LogFloat)] expectedFeaturesIn crf xs = zxF `par` zxB `pseq` zxF `pseq` concat [expectedOn k | k <- [0 .. V.length xs - 1] ] where expectedOn = expectedFeaturesOn crf alpha beta xs alpha = forward crf xs beta = backward crf xs zxF = zxAlpha xs alpha zxB = zxBeta beta