TokenizedDataLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.fillers;
using System.IO;
using MyCaffe.db.image;
using MyCaffe.param.gpt;
using System.Net;
 
namespace MyCaffe.layers.gpt
{
    public class TokenizedDataLayer<T> : Layer<T>
    {
        CancelEvent m_evtCancel;
        InputData m_data;
        Blob<T> m_blobX = null;
        Blob<T> m_blobY = null;
        Random m_random = new Random();
        Layer<T> m_softmax = null;
        bool m_bEnableEos = false;
        bool m_bEnableBos = false;
 
        public TokenizedDataLayer(CudaDnn<T> cuda, Log log, LayerParameter p, IXDatabaseBase db, CancelEvent evtCancel)
            : base(cuda, log, p)
        {
            m_evtCancel = evtCancel;
            m_type = LayerParameter.LayerType.TOKENIZED_DATA;
 
            m_blobX = new Blob<T>(m_cuda, m_log);
        }
 
        protected override void dispose()
        {
            dispose(ref m_blobY);
            dispose(ref m_blobX);
            dispose(ref m_softmax);
 
            base.dispose();
        }
 
        public override int MaxBottomBlobs
        {
            get { return (m_phase == Phase.RUN) ? 1 : 0; }
        }
 
        public override int MinBottomBlobs
        {
            get { return 0; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return (m_phase == Phase.RUN) ? 2 : 3; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nBlockSize = (int)m_param.tokenized_data_param.block_size;
 
            switch (m_param.tokenized_data_param.input_type)
            {
                case TokenizedDataParameter.INPUT_TYPE.TEXT_FILE:
                    m_data = new TextInputData(m_param.tokenized_data_param.source, m_param.tokenized_data_param.vocabulary_type, m_param.tokenized_data_param.seed, m_param.tokenized_data_param.debug_index_file, m_param.phase);
                    break;
 
                default:
                    throw new Exception("Unknown input type '" + m_param.tokenized_data_param.input_type.ToString() + "'");
            }
 
            Reshape(colBottom, colTop);
 
            Blob<T> blobPos = colTop[1];
            List<int> rgShape = Utility.Clone<int>(colTop[1].shape());
            rgShape[1] = nBlockSize;
            blobPos.Reshape(rgShape);
 
            // Set the position data = 0, 1, 2, 3, ... block_size-1
            float[] rgPos = new float[nBlockSize];
            for (int i = 0; i < nBlockSize; i++)
            {
                rgPos[i] = i;
            }
 
            blobPos.mutable_cpu_data = convert(rgPos);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_phase == Phase.RUN)
            {
                int nBatchSize = colBottom[0].num;
                int nBlockSize = (int)m_param.tokenized_data_param.block_size;
                int nTokenSize = (int)m_data.TokenSize;
 
                Blob<T> blobData = colTop[0];
                Blob<T> blobPos = colTop[1];
                Blob<T> blobTarget = null;
                
                if (colTop.Count > 2)
                    blobTarget = colTop[2];
 
                int nCount = 3;
                if (nTokenSize == 1)
                    nCount = 2;
                int[] rgShape = new int[nCount];
 
                blobData.SetParameter("vocab_size", m_data.VocabularySize);
 
                int nC = colBottom[0].channels;
                if (nC > nBlockSize)
                    throw new Exception("The bottom input channel count cannot exceed the block_size=" + nBlockSize.ToString());
 
                // reshape for single characters (each character is an index into the vocab vector)
                rgShape[0] = nBatchSize;
                rgShape[1] = nC;
                if (rgShape.Length > 2)
                    rgShape[2] = nTokenSize;
 
                blobData.Reshape(rgShape);
                if (blobTarget != null)
                    blobTarget.Reshape(rgShape);
 
                if (blobPos.count() < nBlockSize)
                {
                    rgShape[0] = 1;
                    if (rgShape.Length > 2)
                        rgShape[2] = 1;
 
                    blobPos.Reshape(rgShape);
                }
 
                rgShape[0] = nBatchSize;
                rgShape[1] = nC;
                blobPos.Reshape(rgShape);
 
                m_blobX.Reshape(1, nBlockSize, 1, 1);
            }
            else
            {
                m_log.CHECK_EQ(colBottom.Count, 0, "Data Layer takes no input blobs.");
                m_log.CHECK_EQ(colTop.Count, 3, "The TokenizedDataLayer requires 3 top blobs.");
 
                int nBatchSize = (int)m_param.tokenized_data_param.batch_size;
                int nBlockSize = (int)m_param.tokenized_data_param.block_size;
                int nTokenSize = (int)m_data.TokenSize;
 
                Blob<T> blobData = colTop[0];
                Blob<T> blobPos = colTop[1];
                Blob<T> blobTarget = colTop[2];
 
                int nCount = 3;
                if (nTokenSize == 1)
                    nCount = 2;
                int[] rgShape = new int[nCount];
 
                blobData.SetParameter("vocab_size", m_data.VocabularySize);
                // reshape for single characters (each character is an index into the vocab vector)
                rgShape[0] = nBatchSize;
                rgShape[1] = nBlockSize;
                if (rgShape.Length > 2)
                    rgShape[2] = nTokenSize;
 
                blobData.Reshape(rgShape);
                blobTarget.Reshape(rgShape);
 
                rgShape[0] = 1;
                if (rgShape.Length > 2)
                    rgShape[2] = 1;
                blobPos.Reshape(rgShape);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_phase == Phase.RUN)
            {
                m_log.CHECK_EQ(colBottom.Count, 1, "There must be one input blob when running.");
                colTop[0].CopyFrom(colBottom[0]);
                // Top[1] should already have pos data in it.
                // There is no Top[2] target data when running.
            }
            else
            {
                int[] rgnIdx;
                Tuple<float[], float[]> data = m_data.GetData((int)m_param.tokenized_data_param.batch_size, (int)m_param.tokenized_data_param.block_size, null, out rgnIdx);
 
                colTop[0].mutable_cpu_data = convert(data.Item1);
                if (colTop.Count > 2)
                    colTop[2].mutable_cpu_data = convert(data.Item2);
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
        }
 
        public override bool SupportsPreProcessing
        {
            get { return true; }
        }
 
        public override bool SupportsPostProcessingLogits
        {
            get { return true; }
        }
 
        public List<int> Tokenize(string str, bool bAddBos, bool bAddEos)
        {
            m_bEnableBos = bAddBos;
            m_bEnableEos = bAddEos;
            return m_data.Tokenize(str, bAddBos, bAddEos);
        }
 
        public string Detokenize(float[] rg, int nStartIdx, int nCount, bool bIgnoreBos = true, bool bIgnoreEos = true)
        {
            return m_data.Detokenize(rg, nStartIdx, nCount, bIgnoreBos, bIgnoreEos);
        }
 
        public override BlobCollection<T> PreProcessInput(PropertySet customInput, out int nSeqLen, BlobCollection<T> colBottom = null)
        {
            nSeqLen = (int)m_param.tokenized_data_param.block_size;
 
            Blob<T> blobIdx = new Blob<T>(m_cuda, m_log, false);
 
            string strInput = customInput.GetProperty("InputData");
            if (string.IsNullOrEmpty(strInput))
                throw new Exception("Could not find 'InputData' property!");
 
            int[] rgShape = new int[2];
            rgShape[0] = 1;
            rgShape[1] = strInput.Length;
 
            blobIdx.Reshape(rgShape);
 
            List<int> rgTokens = m_data.Tokenize(strInput, false, false);
            float[] rgInput = new float[rgTokens.Count];
 
            for (int i = 0; i < strInput.Length; i++)
            {
                rgInput[i] = rgTokens[i];
            }
 
            blobIdx.mutable_cpu_data = convert(rgInput);
 
            return new BlobCollection<T>() { blobIdx };
        }
 
        public override bool PreProcessInput(string str, int? nTokIdx, BlobCollection<T> colBottom = null)
        {
            if (nTokIdx.HasValue && m_bEnableEos && nTokIdx.Value == (int)SPECIAL_TOKENS.EOS)
                return false;
 
            List<float> rgTok = convertF(colBottom[0].mutable_cpu_data).ToList();
 
            rgTok.Add(nTokIdx.Value);
            if (rgTok.Count > m_param.tokenized_data_param.block_size)
                rgTok.RemoveAt(0);
 
            List<int> rgShape = Utility.Clone<int>(colBottom[0].shape());
            rgShape[1] = rgTok.Count;
            colBottom[0].Reshape(rgShape);
 
            colBottom[0].mutable_cpu_data = convert(rgTok.ToArray());
 
            return true;
        }
 
        public override List<Tuple<string, int, double>> PostProcessLogitsOutput(int nCurIdx, Blob<T> blobLogits, Layer<T> softmax, int nAxis, int nK = 1)
        {
            float[] rgData = convertF(blobLogits.mutable_cpu_data);
            int nVocabCount = blobLogits.count(nAxis);
            float[] rgLogits = new float[nVocabCount];
            int nIdxStart = blobLogits.count() - nVocabCount;
            Dictionary<int, float> rgTopK = new Dictionary<int, float>();
 
            for (int i = nIdxStart; i < blobLogits.count(); i++)
            {
                float fVal = rgData[i];
                rgTopK.Add(i - nIdxStart, fVal);
 
                if (rgTopK.Count > nK)
                {
                    float fMin = float.MaxValue;
                    int nMinIdx = -1;
 
                    foreach (KeyValuePair<int, float> kv in rgTopK)
                    {
                        if (kv.Value < fMin)
                        {
                            fMin = kv.Value;
                            nMinIdx = kv.Key;
                        }
                    }
 
                    rgTopK.Remove(nMinIdx);
                }
            }
 
            for (int i = 0; i < rgLogits.Count(); i++)
            {
                if (rgTopK.ContainsKey(i))
                    rgLogits[i] = rgTopK[i];
                else
                    rgLogits[i] = -float.MaxValue;
            }
 
            if (m_blobX == null)
                m_blobX = new Blob<T>(m_cuda, m_log, false);
            if (m_blobY == null)
                m_blobY = new Blob<T>(m_cuda, m_log, false);
 
            m_blobX.Reshape(1, 1, nVocabCount, 1);
            m_blobX.mutable_cpu_data = convert(rgLogits);
 
            BlobCollection<T> colBottom = new BlobCollection<T>() { m_blobX };
            BlobCollection<T> colTop = new BlobCollection<T>() { m_blobY };
            if (softmax == null)
            {
                if (m_softmax == null)
                {
                    LayerParameter softmax_param = new LayerParameter(LayerParameter.LayerType.SOFTMAX);
                    softmax_param.softmax_param.axis = nAxis;
                    m_softmax = Layer<T>.Create(m_cuda, m_log, softmax_param, null);
                    m_softmax.Setup(colBottom, colTop);
                }
 
                softmax = m_softmax;
            }
            softmax.Forward(colBottom, colTop);
 
            float[] rgProb = convertF(m_blobY.mutable_cpu_data);
            int nCharIdx = (m_param.tokenized_data_param.sample_method == TokenizedDataParameter.SAMPLE_METHOD.PROBABILITY) ? sample(rgProb) : argmax(rgProb);
 
            string str = "";
            str += m_data.Detokenize(nCharIdx, true, true);
 
            return new List<Tuple<string, int, double>>() { new Tuple<string, int, double>(str, nCharIdx, 0) };
        }
 
        private int argmax(float[] rgData)
        {
            int nMaxIdx = 0;
            float fMax = rgData[0];
 
            for (int i = 1; i < rgData.Length; i++)
            {
                if (rgData[i] > fMax)
                {
                    fMax = rgData[i];
                    nMaxIdx = i;
                }
            }
 
            return nMaxIdx;
        }
 
        private int sample(float[] rgData)
        {
            float fTotal = 0;
            float fRand = (float)m_random.NextDouble();
 
            for (int i = 0; i < rgData.Length; i++)
            {
                fTotal += rgData[i];
 
                if (fTotal >= fRand)
                    return i;
            }
 
            return rgData.Length - 1;
        }
    }
 
    public class TextInputData : InputData
    {
        string m_strData;
        IVocabulary m_vocab;
        string m_strDebugIndexFile;
        List<int> m_rgDebugIdx = null;
        int m_nDebugIdx = 0;
        float[] m_rgData = null;
        float[] m_rgTgt = null;
        Phase m_phase;
 
 
        public TextInputData(string strSrc, TokenizedDataParameter.VOCABULARY_TYPE vocabType = TokenizedDataParameter.VOCABULARY_TYPE.CHARACTER, int? nRandomSeed = null, string strDebugIndexFile = null, Phase phase = Phase.NONE) : base(nRandomSeed)
        {
            string strProgData = Environment.GetFolderPath(Environment.SpecialFolder.CommonApplicationData);
            strSrc = Utility.ReplaceMacro(strSrc, "$ProgramData$", strProgData);
 
            m_phase = phase;
            m_strData = File.ReadAllText(strSrc);
 
            if (File.Exists(strDebugIndexFile))
            {
                m_strDebugIndexFile = Utility.ReplaceMacro(strDebugIndexFile, "$ProgramData$", strProgData);
                m_rgDebugIdx = new List<int>();
                string[] rgLines = File.ReadAllLines(strDebugIndexFile);
                foreach (string strLine in rgLines)
                {
                    if (strLine.StartsWith("idx = "))
                    {
                        string strIdx = strLine.Substring(6).Trim(' ', '\t', '\n', '\r');
                        m_rgDebugIdx.Add(int.Parse(strIdx));
                    }
                }
            }
 
            if (vocabType == TokenizedDataParameter.VOCABULARY_TYPE.WORD)
                m_vocab = new VocabularyWord(m_random, false, false);
            else
               m_vocab = new VocabularyCharacter(m_random, false, false, false);
            
            m_vocab.BuildFromString(m_strData);
        }
 
        public override List<string> RawData
        {
            get
            {
                return new List<string>() { m_strData };
            }
        }
 
        public override uint TokenSize
        {
            get { return 1; }
        }
 
        public override uint VocabularySize
        {
            get { return (uint)m_vocab.Count; }
        }
 
        public override bool GetDataAvailabilityAt(int nIdx, bool bIncludeSrc, bool bIncludeTrg)
        {
            return true;
        }
 
        public override Tuple<float[], float[]> GetData(int nBatchSize, int nBlockSize, InputData trgData, out int[] rgnIdx)
        {
            int nSize = nBatchSize * nBlockSize;
 
            rgnIdx = new int[nBatchSize];
 
            if (m_rgData == null || m_rgData.Length != nSize)
                m_rgData = new float[nSize];
 
            if (m_rgTgt == null || m_rgTgt.Length != nSize)
                m_rgTgt = new float[nSize];
 
            for (int i = 0; i < nBatchSize; i++)
            {
                int nMax = m_strData.Count() - (nBlockSize + 1);
                int nDataIdx = m_random.Next(nMax);
                int nDstIdx = i * nBlockSize;
 
                rgnIdx[i] = nDataIdx;
 
                if (m_rgDebugIdx != null)
                {
                    nDataIdx = m_rgDebugIdx[m_nDebugIdx];
                    m_nDebugIdx++;
 
                    if (m_nDebugIdx >= m_rgDebugIdx.Count)
                        m_nDebugIdx = 0;
                }
 
                List<int> rgTokens = new List<int>();
                List<int> rgLastTokens;
                int nIdx = 0;
                
                while (rgTokens.Count < nBlockSize + 1)
                {
                    rgLastTokens = m_vocab.Tokenize(m_strData[nDataIdx + nIdx].ToString());
                    if (rgLastTokens.Count > 0)
                        rgTokens.AddRange(rgLastTokens);
 
                    nIdx++;
                }
 
                Array.Copy(rgTokens.ToArray(), 0, m_rgData, nDstIdx, nBlockSize);
                rgTokens.RemoveAt(0);
                Array.Copy(rgTokens.ToArray(), 0, m_rgTgt, nDstIdx, nBlockSize);
            }
 
            return new Tuple<float[], float[]>(m_rgData, m_rgTgt);
        }
 
        public override Tuple<float[], float[]> GetDataAt(int nBatchSize, int nBlockSize, int[] rgnIdx)
        {
            throw new NotImplementedException();
        }
 
        public override List<int> Tokenize(string str, bool bAddBos, bool bAddEos)
        {
            return m_vocab.Tokenize(str, bAddBos, bAddEos).ToList();
        }
        
        public override string Detokenize(int nTokIdx, bool bIgnoreBos, bool bIgnoreEos)
        {
            return m_vocab.Detokenize(nTokIdx, bIgnoreBos, bIgnoreEos);
        }
 
        public override string Detokenize(float[] rgfTokIdx, int nStartIdx, int nCount, bool bIgnoreBos, bool bIgnoreEos)
        {
            string str = "";
            for (int i = nStartIdx; i < nStartIdx + nCount; i++)
            {
                string strItem = m_vocab.Detokenize((int)rgfTokIdx[i], bIgnoreBos, bIgnoreEos);
                if (string.IsNullOrEmpty(strItem))
                    break;
 
                str += strItem;
            }
 
            return str;
        }
 
        public override char BOS
        { 
            get { return m_vocab.BOS; }
        }
 
        public override char EOS
        {
            get { return m_vocab.EOS; }
        }
    }
}