mirror of
https://github.com/ilyakooo0/grenade.git
synced 2024-09-11 14:45:57 +03:00
188 lines
6.8 KiB
Haskell
188 lines
6.8 KiB
Haskell
{-# LANGUAGE BangPatterns #-}
|
|
{-# LANGUAGE CPP #-}
|
|
{-# LANGUAGE DataKinds #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
{-# LANGUAGE TypeOperators #-}
|
|
{-# LANGUAGE TupleSections #-}
|
|
{-# LANGUAGE TypeFamilies #-}
|
|
{-# LANGUAGE LambdaCase #-}
|
|
|
|
import Control.Monad.Random
|
|
import Control.Monad.Trans.Except
|
|
|
|
import Data.Char ( isUpper, toUpper, toLower )
|
|
import Data.List ( foldl' )
|
|
import Data.Maybe ( fromMaybe )
|
|
|
|
|
|
import qualified Data.Vector as V
|
|
import Data.Vector ( Vector )
|
|
|
|
import qualified Data.Map as M
|
|
|
|
import qualified Data.ByteString as B
|
|
import Data.Serialize
|
|
|
|
import GHC.TypeLits
|
|
|
|
import Numeric.LinearAlgebra.Static ( konst )
|
|
|
|
import Options.Applicative
|
|
|
|
import Grenade
|
|
import Grenade.Recurrent
|
|
import Grenade.Utils.OneHot
|
|
|
|
import System.IO.Unsafe ( unsafeInterleaveIO )
|
|
import Data.Proxy
|
|
import Prelude.Singletons
|
|
|
|
-- The defininition for our natural language recurrent network.
|
|
-- This network is able to learn and generate simple words in
|
|
-- about an hour.
|
|
--
|
|
-- Grab the input from
|
|
-- https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
|
|
--
|
|
-- This is a first class recurrent net.
|
|
--
|
|
-- The F and R types are tagging types to ensure that the runner and
|
|
-- creation function know how to treat the layers.
|
|
--
|
|
-- As an example, here's a short sequence generated.
|
|
--
|
|
-- > KING RICHARD III:
|
|
-- > And as the heaven her his words, we the son, I show sand stape but the lament to shall were the sons with a strend
|
|
|
|
type F = FeedForward
|
|
type R = Recurrent
|
|
|
|
-- The definition of our network
|
|
type Shakespeare = RecurrentNetwork '[ R (LSTM 40 80), R (LSTM 80 40), F (FullyConnected 40 40), F Logit]
|
|
'[ 'D1 40, 'D1 80, 'D1 40, 'D1 40, 'D1 40 ]
|
|
|
|
-- The definition of the "sideways" input, which the network is fed recurrently.
|
|
type Shakespearian = RecurrentInputs '[ R (LSTM 40 80), R (LSTM 80 40), F (FullyConnected 40 40), F Logit]
|
|
|
|
randomNet :: MonadRandom m => m Shakespeare
|
|
randomNet = randomRecurrent
|
|
|
|
-- | Load the data files and prepare a map of characters to a compressed int representation.
|
|
loadShakespeare :: FilePath -> ExceptT String IO (Vector Int, M.Map Char Int, Vector Char)
|
|
loadShakespeare path = do
|
|
contents <- lift $ readFile path
|
|
let annotated = annotateCapitals contents
|
|
(m,cs) <- ExceptT . return $ hotMap (Proxy :: Proxy 40) annotated
|
|
hot <- ExceptT . return . note "Couldn't generate hot values" $ traverse (`M.lookup` m) annotated
|
|
return (V.fromList hot, m, cs)
|
|
|
|
trainSlice :: LearningParameters -> Shakespeare -> Shakespearian -> Vector Int -> Int -> Int -> (Shakespeare, Shakespearian)
|
|
trainSlice !lrate !net !recIns input offset size =
|
|
let e = fmap (x . oneHot) . V.toList $ V.slice offset size input
|
|
in case reverse e of
|
|
(o : l : xs) ->
|
|
let examples = reverse $ (l, Just o) : ((,Nothing) <$> xs)
|
|
in trainRecurrent lrate net recIns examples
|
|
_ -> error "Not enough input"
|
|
where
|
|
x = fromMaybe (error "Hot variable didn't fit.")
|
|
|
|
runShakespeare :: ShakespeareOpts -> ExceptT String IO ()
|
|
runShakespeare opts = do
|
|
(shakespeare, oneHotMap, oneHotDictionary) <- loadShakespeare $ trainingFile opts
|
|
(net0, i0) <- lift $
|
|
case loadPath opts of
|
|
Just loadFile -> netLoad loadFile
|
|
Nothing -> (,0) <$> randomNet
|
|
|
|
(trained, bestInput) <- lift $ foldM (\(!net, !io) size -> do
|
|
xs <- take (iterations opts `div` 10) <$> getRandomRs (0, length shakespeare - size - 1)
|
|
let (!trained, !bestInput) = foldl' (\(!n, !i) offset -> trainSlice (rate opts) n i shakespeare offset size) (net, io) xs
|
|
results <- take 1000 <$> generateParagraph trained bestInput (temperature opts) oneHotMap oneHotDictionary ( S1D $ konst 0)
|
|
putStrLn ("TRAINING STEP WITH SIZE: " ++ show size)
|
|
putStrLn (unAnnotateCapitals results)
|
|
return (trained, bestInput)
|
|
) (net0, i0) $ replicate 10 (sequenceSize opts)
|
|
|
|
case savePath opts of
|
|
Just saveFile -> lift . B.writeFile saveFile $ runPut (put trained >> put bestInput)
|
|
Nothing -> return ()
|
|
|
|
generateParagraph :: forall layers shapes n a. (Last shapes ~ 'D1 n, Head shapes ~ 'D1 n, KnownNat n, Ord a)
|
|
=> RecurrentNetwork layers shapes
|
|
-> RecurrentInputs layers
|
|
-> Double
|
|
-> M.Map a Int
|
|
-> Vector a
|
|
-> S ('D1 n)
|
|
-> IO [a]
|
|
generateParagraph n s temp hotmap hotdict =
|
|
go s
|
|
where
|
|
go x y =
|
|
do let (_, ns, o) = runRecurrent n x y
|
|
un <- sample temp hotdict o
|
|
Just re <- return $ makeHot hotmap un
|
|
rest <- unsafeInterleaveIO $ go ns re
|
|
return (un : rest)
|
|
|
|
data ShakespeareOpts = ShakespeareOpts {
|
|
trainingFile :: FilePath
|
|
, iterations :: Int
|
|
, rate :: LearningParameters
|
|
, sequenceSize :: Int
|
|
, temperature :: Double
|
|
, loadPath :: Maybe FilePath
|
|
, savePath :: Maybe FilePath
|
|
}
|
|
|
|
shakespeare' :: Parser ShakespeareOpts
|
|
shakespeare' = ShakespeareOpts <$> argument str (metavar "TRAIN")
|
|
<*> option auto (long "examples" <> short 'e' <> value 1000000)
|
|
<*> (LearningParameters
|
|
<$> option auto (long "train_rate" <> short 'r' <> value 0.01)
|
|
<*> option auto (long "momentum" <> value 0.95)
|
|
<*> option auto (long "l2" <> value 0.000001)
|
|
)
|
|
<*> option auto (long "sequence-length" <> short 's' <> value 50)
|
|
<*> option auto (long "temperature" <> short 't' <> value 0.4)
|
|
<*> optional (strOption (long "load"))
|
|
<*> optional (strOption (long "save"))
|
|
|
|
main :: IO ()
|
|
main = do
|
|
shopts <- execParser (info (shakespeare' <**> helper) idm)
|
|
res <- runExceptT $ runShakespeare shopts
|
|
case res of
|
|
Right () -> pure ()
|
|
Left err -> putStrLn err
|
|
|
|
|
|
netLoad :: FilePath -> IO (Shakespeare, Shakespearian)
|
|
netLoad modelPath = do
|
|
modelData <- B.readFile modelPath
|
|
either fail return $ runGet get modelData
|
|
|
|
-- Replace capitals with an annotation and the lower case letter
|
|
-- http://fastml.com/one-weird-trick-for-training-char-rnns/
|
|
annotateCapitals :: String -> String
|
|
annotateCapitals (x : rest)
|
|
| isUpper x
|
|
= '^' : toLower x : annotateCapitals rest
|
|
| otherwise
|
|
= x : annotateCapitals rest
|
|
annotateCapitals []
|
|
= []
|
|
|
|
unAnnotateCapitals :: String -> String
|
|
unAnnotateCapitals ('^' : x : rest)
|
|
= toUpper x : unAnnotateCapitals rest
|
|
unAnnotateCapitals (x : rest)
|
|
= x : unAnnotateCapitals rest
|
|
unAnnotateCapitals []
|
|
= []
|
|
|
|
-- | Tag the 'Nothing' value of a 'Maybe'
|
|
note :: a -> Maybe b -> Either a b
|
|
note a = maybe (Left a) Right
|