colada-0.0.1: colada.hs

{-# LANGUAGE FlexibleInstances , DeriveDataTypeable 
 , TemplateHaskell , OverloadedStrings
 #-}
module Main
where       
import qualified Data.Text.Lazy.IO as Text
import qualified Data.Text.Lazy as Text
import qualified Data.Text.Lazy.Builder as Text
import qualified Data.Text.Lazy.Builder.Int as Text
import qualified Data.ByteString as BS
import qualified Data.Serialize as Serialize
import qualified Data.List as List
import qualified Data.Vector.Generic as V

import qualified System.Environment as Env
import System.Console.CmdArgs.Explicit
import qualified Data.Label as L
import qualified Data.Label.Maybe as M
import Prelude hiding ((.))
import Control.Category ((.))

import qualified NLP.CoNLL as CoNLL
import qualified Colada.WordClass as C

-- Command line parsing

data Program = Help 
             | Learn { _options :: C.Options 
                     , _modelPath :: FilePath }
             | Predict { _modelPath :: FilePath }
             | Label { _modelPath :: FilePath }
             deriving (Show)
$(L.mkLabels [''Program])                                       

help :: Mode Program
help = 
  mode "help" Help "Display help" 
  (flagArg (\ x _ -> 
             Left $ "Unexpected argument " ++ x) "")
  []

predict :: Mode Program
predict = 
  mode "predict" Predict { _modelPath = "model" } "Predict words"
  (flagArg (\x p -> Right $ maybe p id (M.set modelPath x p)) "FILE")
  []

label :: Mode Program
label =
  mode "label" Label { _modelPath = "model" } "Label words with classes"
  (flagArg (\x p -> Right $ maybe p id (M.set modelPath x p)) "FILE")
  []
  
learn :: Mode Program
learn = 
  let setOption field x p =   
          fmap (maybe p id . flip (M.set (field . options)) p)
        . safeRead 
        $ x 
      safeRead :: Read b => String -> Either String b
      safeRead x = 
        case reads x of
          [(a,"")] -> Right a
          _        -> Left $ "Couldn't parse " ++ show x
  in mode "learn" Learn { _options = C.defaultOptions 
                        , _modelPath = "model" } "Learn word classes"
     (flagArg (\x p -> Right $ maybe p id (M.set modelPath x p))
      "FILE")
        [ flagReq ["features"]  
            (\x p -> case x of 
                "unigram" -> Right . maybe p id 
                             $ M.set (C.featIds . options) [-1,1] p
                "bigram"  -> Right . maybe p id 
                             $ M.set (C.featIds . options) [-12,12] p
                _         -> Left $ "Unknown feature specification " ++ x)
            "(unigram|bigram)" "Feature specification"
          
        , flagReq ["topic-num"] (setOption C.topicNum)
            "NAT" "Number of topics K" 
        
        , flagReq ["alphasum"] (setOption C.alphasum)
            "FLOAT" "Parameter alpha * K"
        
        , flagReq ["beta"] (setOption C.beta)  
            "FLOAT" "Parameter beta"
          
        , flagReq ["passes"] (setOption C.passes)
            "NAT" "Passes per batch"
          
        , flagReq ["repeats"] (setOption C.repeats)
            "NAT" "Repeats per sentence"
          
        , flagReq ["batch-size"] (setOption C.batchSize)
            "NAT" "Sentences per batch"
          
        , flagReq ["seed"] (setOption C.seed)
            "NAT" "Random seed"
        ]
        
  
program :: Mode Program                     
program = modes "colada" Help "Word class learning" 
          [learn, predict, label, help] 


-- Run the program

main :: IO ()
main = do
  args <- Env.getArgs
  let opts = processValue program args
  case opts of
    Help -> print $ helpText [] HelpFormatDefault program
    Predict { _modelPath = p } -> do
      -- FIXME: use Data.Text.Builder instead of converting to Lists
      let format s = {-# SCC "format" #-}
                   Text.unlines 
                     [ Text.concat . List.intersperse "," . map snd . V.toList 
                       $ ws 
                     | ws <-  V.toList s ]
      m <- parseModel p
      ss <- CoNLL.parse `fmap` Text.getContents
      Text.putStr . Text.unlines . map (format . C.predict m) $ ss
    Label { _modelPath = p } -> do
      let format s = Text.unlines 
                     . V.toList 
                     . V.map (Text.toLazyText . Text.decimal . V.maxIndex) 
                     $ s
      m <- parseModel p
      ss <- CoNLL.parse `fmap` Text.getContents
      Text.putStr . Text.unlines . map (format . C.label m) $ ss  
    Learn { _options = o , _modelPath = p } -> do
      ss <- CoNLL.parse `fmap` Text.getContents
      let m = C.run o ss
      Text.putStr . C.summary $ m
      BS.writeFile p . Serialize.encode $ m

parseModel :: FilePath -> IO C.WordClass
parseModel p = do
  (either (\err -> error $ "Error reading model " ++ err) id
   . Serialize.decode)
  `fmap` BS.readFile p