-- ---
-- jupyter:
--   jupytext:
--     text_representation:
--       extension: .hs
--       format_name: light
--       format_version: '1.5'
--       jupytext_version: 1.14.4
--   kernelspec:
--     display_name: Haskell - Tensorflow
--     language: haskell
--     name: ihaskell_tensorflow
-- ---

-- # Tensorflow example with datasets
--
-- _example after https://mmhaskell.com/blog/2017/8/21/digging-in-deep-solving-a-real-problem-with-haskell-tensor-flow_
--
--
-- here are various inputs that we are going to need throughout the code:

-- +
:e DeriveGeneric
:e OverloadedLists
:e TypeFamilies

import Numeric.Datasets.Iris (iris, Iris(..), IrisClass(..))
import Control.Monad (forM_, when)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString.Lazy.Char8 (pack)
import Data.Csv (FromRecord, decode, HasHeader(..))
import Data.Int (Int64)
import Data.Vector (Vector, length, fromList, (!))
import GHC.Generics (Generic)
import System.Random.Shuffle (shuffleM)

import TensorFlow.Core (TensorData, Session, Build, render, runWithFeeds, feed, unScalar, build,
                        Tensor, encodeTensorData)
import TensorFlow.Minimize (minimizeWith, adam', AdamConfig(..))
import TensorFlow.Ops (placeholder, truncatedNormal, add, matMul, relu,
                      argMax, scalar, cast, oneHot, reduceMean, softmaxCrossEntropyWithLogits, 
                      equal, vector)
import TensorFlow.Session (runSession)
import TensorFlow.Types (Shape(..))
import TensorFlow.Variable (readValue, initializedVariable, Variable)
-- -

-- our input data comes from the *Numeric.Dataset.Iris* module. It is a list of data samples that are represented by the type *Iris*:

:t Iris
:t iris

-- for Tensorflow, we need to know the input data size, the number of labels and the number of features. The labels have to be converted to a numberical value.

nlabels = 3 :: Int64
nfields = 4 :: Int64
irisClassToInt64 :: IrisClass -> Int64
irisClassToInt64 x = case x of
                Setosa -> -1
                Versicolor -> 0
                Virginica -> 1

-- we can now convert the dataset into the required form.

convertRecordsToTensorData :: Vector Iris -> (TensorData Float, TensorData Int64)
convertRecordsToTensorData records = (input, output)
  where
    numRecords = Data.Vector.length records
    input = encodeTensorData [fromIntegral numRecords, nfields] (fromList $ concatMap recordToInputs records)
    output = encodeTensorData [fromIntegral numRecords] ((irisClassToInt64 . irisClass) <$> records)
    recordToInputs :: Iris -> [Float]
    recordToInputs rec = realToFrac <$> [sepalLength rec, sepalWidth rec, petalLength rec, petalWidth rec]

-- +
sampleSize :: Int
sampleSize = 100

chooseRandomRecords :: Vector Iris -> IO (Vector Iris)
chooseRandomRecords records = do
  let numRecords = Data.Vector.length records
  chosenIndices <- take sampleSize <$> shuffleM [0..(numRecords - 1)]
  return $ fromList $ map (records !) chosenIndices
:t chooseRandomRecords
-- -

data Model = Model
  { train :: TensorData Float -- Training input
          -> TensorData Int64 -- Training output
          -> Session ()
  , errorRate :: TensorData Float -- Test input
          -> TensorData Int64 -- Test output
          -> Session Float
  }
:t Model

buildNNLayer :: Int64 -> Int64 -> Tensor v Float
             -> Build (Variable Float, Variable Float, Tensor Build Float)
buildNNLayer inputSize outputSize input = do
  weights <- truncatedNormal (vector [inputSize, outputSize]) >>= initializedVariable
  bias <- truncatedNormal (vector [outputSize]) >>= initializedVariable
  let results = (input `matMul` readValue weights) `add` readValue bias
  return (weights, bias, results)
:t buildNNLayer

createModel :: Build Model
createModel = do
  let batchSize = -1 -- Allows variable sized batches
  let numHiddenUnits = 50
  
  inputs <- placeholder [batchSize, nfields]
  outputs <- placeholder [batchSize]
  
  -- first layer (input -> hidden)
  (hiddenWeights, hiddenBiases, hiddenResults) <- buildNNLayer nfields numHiddenUnits inputs
  let rectifiedHiddenResults = relu hiddenResults
  
  -- output layer (hidden -> output)
  (finalWeights, finalBiases, finalResults) <- buildNNLayer numHiddenUnits nlabels rectifiedHiddenResults
  actualOutput <- render $ argMax finalResults (scalar (1 :: Int64))
  
  let correctPredictions = equal actualOutput outputs
  errorRate_ <- render $ 1 - (reduceMean (cast correctPredictions))
  let outputVectors = oneHot outputs (fromIntegral nlabels) 1 0
  
  let loss = reduceMean $ fst $ softmaxCrossEntropyWithLogits finalResults outputVectors
  let params = [hiddenWeights, hiddenBiases, finalWeights, finalBiases]
  
  let adamConfig = AdamConfig 1e-3 0.9 0.999 1e-8
  train_ <- minimizeWith (adam' adamConfig) loss params
  
  return $ Model
      { train = \inputFeed outputFeed ->
          runWithFeeds
            [ feed inputs inputFeed
            , feed outputs outputFeed
            ]
            train_
      , errorRate = \inputFeed outputFeed -> unScalar <$>
          runWithFeeds
        [ feed inputs inputFeed
        , feed outputs outputFeed
        ]
        errorRate_
      }

-- +
runIris = runSession $ do
  model <- build createModel

  forM_ ([0..10000] :: [Int]) $ \i -> do
    trainingSample <- liftIO $ chooseRandomRecords $ Data.Vector.fromList iris
    let (trainingInputs, trainingOutputs) = convertRecordsToTensorData trainingSample
    let (allInputs, allOutputs) = convertRecordsToTensorData $ Data.Vector.fromList iris
    train model trainingInputs trainingOutputs
    when (i `mod` 2000 == 0) $ do
      err <- errorRate model allInputs allOutputs
      liftIO $ putStrLn $ "Current training error " ++ show err

runIris

StarcoderdataJupyterScriptsDedupFiltered