SHARE
TWEET

Untitled

a guest Dec 16th, 2018 64 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. module Markov where
  2.  
  3. import Control.Monad (guard)
  4. import Data.Monoid
  5. import System.Random (StdGen,mkStdGen,getStdGen,setStdGen,randomR,randomRIO)
  6. import System.IO.Unsafe (unsafeInterleaveIO)
  7. import qualified Data.Foldable as F
  8. import qualified Data.Map as Map
  9. import qualified Data.Set as Set
  10. import qualified Data.Text.Lazy as L
  11. import qualified Data.Text.Lazy.IO as L
  12.  
  13. import Debug.Trace
  14.  
  15.  
  16. data Markov m a = Markov
  17.   { mStart :: Set.Set a
  18.     -- ^ Set of starting atoms
  19.   , mNext  :: Map.Map a (Map.Map a m)
  20.     -- ^ Map of transitions between these two atoms
  21.   } deriving (Show)
  22.  
  23. instance (Monoid m, Ord a) => Monoid (Markov m a) where
  24.   mempty  = Markov
  25.     { mStart = Set.empty
  26.     , mNext  = Map.empty
  27.     }
  28.  
  29.   mappend l r = Markov
  30.     { mStart = Set.union (mStart l) (mStart r)
  31.     , mNext  = Map.unionWith combine (mNext l) (mNext r)
  32.     }
  33.  
  34. combine :: (Monoid m, Ord a) => Map.Map a m -> Map.Map a m -> Map.Map a m
  35. combine  = Map.unionWith mappend
  36.  
  37.  
  38. -- Learning Mode ---------------------------------------------------------------
  39.  
  40. type MarkovL = Markov (Sum Int)
  41.  
  42. fromList :: Ord a => [a] -> MarkovL a
  43. fromList as0 = case as0 of
  44.   a:rest -> loop (start a mempty) a rest
  45.   _      -> error "fromList: empty input"
  46.   where
  47.   loop m a as = case as of
  48.     b:rest -> loop (next a b m) b rest
  49.     []     -> m
  50.  
  51. -- | Add an atom to the starting set.
  52. start :: Ord a => a -> MarkovL a -> MarkovL a
  53. start a m = m { mStart = Set.insert a (mStart m) }
  54.  
  55. -- | Increment the edge count between two atoms.
  56. next :: Ord a => a -> a -> MarkovL a -> MarkovL a
  57. next a a' m = m
  58.   { mNext = Map.insertWith combine a (Map.singleton a' (Sum 1)) (mNext m)
  59.   }
  60.  
  61.  
  62. -- Compiled Mode ---------------------------------------------------------------
  63.  
  64. type MarkovC = Markov Double
  65.  
  66. compile :: Ord a => MarkovL a -> MarkovC a
  67. compile m = m { mNext = Map.mapWithKey update (mNext m) }
  68.   where
  69.   update k as = Map.map (\p -> fromIntegral (getSum p) / total) as
  70.     where
  71.     total = fromIntegral (getSum (F.fold (Map.elems as)))
  72.  
  73. stepsIO :: Ord a => MarkovC a -> a -> IO [a]
  74. stepsIO s a0 = (a0:) `fmap` loop a0
  75.   where
  76.   loop a = do
  77.     mb <- stepIO s a
  78.     case mb of
  79.       Just a' -> do
  80.         rest <- unsafeInterleaveIO (loop a')
  81.         return (a':rest)
  82.       Nothing -> return []
  83.  
  84. stepsIO' :: Ord a => MarkovC a -> IO [a]
  85. stepsIO' s = do
  86.   let starts = Set.toList (mStart s)
  87.   ix <- randomRIO (0,length starts - 1)
  88.   stepsIO s (starts !! ix)
  89.  
  90. stepIO :: Ord a => MarkovC a -> a -> IO (Maybe a)
  91. stepIO s a = do
  92.   r <- getStdGen
  93.   case step s a r of
  94.     Just (a',r') -> do
  95.       setStdGen r'
  96.       return (Just a')
  97.     Nothing -> return Nothing
  98.  
  99. step :: Ord a => MarkovC a -> a -> StdGen -> Maybe (a,StdGen)
  100. step s a r = do
  101.   ns <- Map.lookup a (mNext s)
  102.   choose (Map.toList ns) r
  103.  
  104. choose :: [(a,Double)] -> StdGen -> Maybe (a,StdGen)
  105. choose as r = do
  106.   guard (not (null as))
  107.   let (val,r') = randomR (0,1) r
  108.       loop p xs = case xs of
  109.         (a,ap):rest | ap > p    -> Just (a,r')
  110.                     | otherwise -> loop (p - ap) rest
  111.         []                      -> Nothing
  112.   loop val as
  113.  
  114.  
  115.  
  116. -- Tests -----------------------------------------------------------------------
  117.  
  118. parse :: FilePath -> IO (MarkovC L.Text)
  119. parse path = do
  120.   input <- L.readFile path
  121.   return $ compile
  122.          $ mconcat
  123.          $ map (fromList . L.words)
  124.          $ filter (not . L.null)
  125.          $ L.lines input
  126.  
  127. sample :: MarkovC L.Text -> IO ()
  128. sample m = do
  129.   ws <- stepsIO' m
  130.   L.putStrLn (L.unwords ws)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top