mirror of
https://github.com/facebook/duckling.git
synced 2024-10-27 14:26:54 +03:00
69 lines
2.1 KiB
Haskell
69 lines
2.1 KiB
Haskell
-- Copyright (c) 2016-present, Facebook, Inc.
|
|
-- All rights reserved.
|
|
--
|
|
-- This source code is licensed under the BSD-style license found in the
|
|
-- LICENSE file in the root directory of this source tree. An additional grant
|
|
-- of patent rights can be found in the PATENTS file in the same directory.
|
|
|
|
|
|
{-# LANGUAGE RecordWildCards #-}
|
|
{-# LANGUAGE NoRebindableSyntax #-}
|
|
|
|
module Duckling.Ranking.Rank
|
|
( rank
|
|
) where
|
|
|
|
import Control.Arrow ((***))
|
|
import Control.Monad (join)
|
|
import qualified Data.HashMap.Strict as HashMap
|
|
import Data.HashSet (HashSet)
|
|
import qualified Data.HashSet as HashSet
|
|
import Data.Maybe
|
|
import qualified Data.Set as Set
|
|
import Prelude
|
|
|
|
import Duckling.Dimensions.Types
|
|
import Duckling.Ranking.Extraction
|
|
import Duckling.Ranking.Types
|
|
import Duckling.Types
|
|
|
|
classify :: Classifier -> BagOfFeatures -> (Class, Double)
|
|
classify Classifier {..} feats = if okScore >= koScore
|
|
then (True, okScore)
|
|
else (False, koScore)
|
|
where
|
|
(okScore, koScore) = join (***) (p feats) (okData, koData)
|
|
p :: BagOfFeatures -> ClassData -> Double
|
|
p feats ClassData{..} =
|
|
prior + HashMap.foldrWithKey (\feat x res ->
|
|
res + fromIntegral x * HashMap.lookupDefault unseen feat likelihoods
|
|
) 0.0 feats
|
|
|
|
score :: Classifiers -> Node -> Double
|
|
score classifiers node@Node {rule = Just rule, ..} =
|
|
case HashMap.lookup rule classifiers of
|
|
Just c -> let feats = extractFeatures node
|
|
in snd (classify c feats) + sum (map (score classifiers) children)
|
|
Nothing -> 0.0
|
|
score _ Node {rule = Nothing} = 0.0
|
|
|
|
-- | Return all superior candidates, as defined by the partial ordering
|
|
winners :: Ord a => [a] -> [a]
|
|
winners xs = filter (\x -> all ((/=) LT . compare x) xs) xs
|
|
|
|
-- | Return a curated list of tokens
|
|
rank
|
|
:: Classifiers
|
|
-> HashSet (Some Dimension)
|
|
-> [ResolvedToken]
|
|
-> [ResolvedToken]
|
|
rank classifiers targets tokens =
|
|
Set.toList . Set.fromList
|
|
. map (\(Candidate token _ _) -> token)
|
|
. winners
|
|
$ map makeCandidate tokens
|
|
where
|
|
makeCandidate :: ResolvedToken -> Candidate
|
|
makeCandidate token@Resolved {node = n@Node {token = Token d _}} =
|
|
Candidate token (score classifiers n) $ HashSet.member (Some d) targets
|