Test batches integration with model.

FastText/src/FastText.jl

@@ -39,6 +39,7 @@ using DataDeps
 using BSON
 using TextAnalysis
 using MLUtils
+using Zygote
 
 
 include("recipes.jl")

FastText/src/encodings/textpreprocessing.jl

@@ -70,7 +70,7 @@ function computevocabulary(data; vocab_size=40000)
     counter = 3
 
     for (k, v) in ordered_dict
-        ordered_dict[k] = counter + 1
+        ordered_dict[k] = counter
         counter = counter + 1
     end
 

FastText/src/models/pretrain_lm.jl

@@ -26,9 +26,9 @@ mutable struct LanguageModel
     layers :: Flux.Chain
 end
 
-function LanguageModel(load_pretrained::Bool=false, vocabpath::String=joinpath(@__DIR__,"vocabs/lm_vocab.csv");embedding_size::Integer=400, hid_lstm_sz::Integer=1150, out_lstm_sz::Integer=embedding_size,
+function LanguageModel(load_pretrained::Bool=false, task::Any = Nothing;embedding_size::Integer=400, hid_lstm_sz::Integer=1150, out_lstm_sz::Integer=embedding_size,
     embed_drop_prob::Float64 = 0.05, in_drop_prob::Float64 = 0.4, hid_drop_prob::Float64 = 0.5, layer_drop_prob::Float64 = 0.3, final_drop_prob::Float64 = 0.3)
-    vocab = (string.(readdlm(vocabpath, ',')))[:, 1]
+    vocab = task.encodings[3].vocab.keys
     de = gpu(DroppedEmbeddings(length(vocab), embedding_size, embed_drop_prob; init = (dims...) -> init_weights(0.1, dims...)))
     lm = LanguageModel(
         vocab,
@@ -45,7 +45,7 @@ function LanguageModel(load_pretrained::Bool=false, vocabpath::String=joinpath(@
             softmax
         )
     )
-    load_pretrained && load_model!(lm, datadep"Pretrained ULMFiT Language Model/ulmfit_lm_en.bson")
+    # load_pretrained && load_model!(lm, datadep"Pretrained ULMFiT Language Model/ulmfit_lm_en.bson")
     return lm
 end
 

FastText/src/models/train_text_classifier.jl

@@ -89,21 +89,23 @@ gen             : data loader, which will give 'X' of the mini-batch in one call
 tracked_steps   : This is the number of tracked time-steps for Truncated Backprop thorugh time,
                   these will be last time-steps for which gradients will be calculated.
 """
-function forward(tc::TextClassifier, gen::Channel, tracked_steps::Integer=32)
+function forward(tc::TextClassifier, batches, tracked_steps::Integer=32)
   	# swiching off tracking
     classifier = tc
-    X = take!(gen)
-    # println("X = $X")
+    # X = take!(gen)
+    X = batches[1][1]
     l = length(X)
     # Truncated Backprop through time
     println("l = $l")
     Zygote.ignore() do
 	for i=1:ceil(l/tracked_steps)-1   # Tracking is swiched off inside this loop
         println("i = $i / $(ceil(l/tracked_steps)-1)")
 	    (i == 1 && l%tracked_steps != 0) ? (last_idx = l%tracked_steps) : (last_idx = tracked_steps)
-	    H = broadcast(x -> indices(x, classifier.vocab, "_unk_"), X[1:last_idx])
+	    # H = broadcast(x -> indices(x, classifier.vocab, "_unk_"), X[1:last_idx])
+        H = X[1:last_idx]
 	    H = classifier.rnn_layers.(H)
 	    X = X[last_idx+1:end]
+        println(length(X))
 	end
 
     println("Start shifting states")
@@ -125,7 +127,8 @@ function forward(tc::TextClassifier, gen::Channel, tracked_steps::Integer=32)
     end
     println("End shifting")
     # last part of the sequecnes in X - Tracking is swiched on
-    H = broadcast(x -> tc.rnn_layers[1](indices(x, classifier.vocab, "_unk_")), X)
+    # H = broadcast(x -> tc.rnn_layers[1](indices(x, classifier.vocab, "_unk_")), X)
+    H = classifier.rnn_layers[1](X[1])
     H = tc.rnn_layers[2:end].(H)
     H = tc.linear_layers(H)
     return H
@@ -144,20 +147,21 @@ classifier    : Instance of TextClassifier
 gen           : 'Channel' [data loader], to give a mini-batch
 tracked_steps : specifies the number of time-steps for which tracking is on
 """
-function loss(classifier::TextClassifier, gen::Channel, tracked_steps::Integer=32)
-    H = forward(classifier, gen, tracked_steps)
-    Y = gpu(take!(gen))
-    l = crossentropy(H, Y)
+function loss(classifier::TextClassifier, batches, tracked_steps::Integer=32)
+    H = forward(classifier, batches, tracked_steps)
+    # Y = gpu(take!(gen))
+    Y = batches[1][2]
+    l = Flux.Losses.crossentropy(H, Y)
     # reset!(classifier.rnn_layers)
     println("Loss = $l")
     return l
 end
 
-function discriminative_step!(layers, classifier::TextClassifier, gen::Channel, tracked_steps::Integer, ηL::Float64, opts::Vector)
+function discriminative_step!(layers, classifier::TextClassifier, batches, tracked_steps::Integer, ηL::Float64, opts::Vector)
     @assert length(opts) == length(layers)
     # Gradient calculation
     println("Start grads")
-    grads = Zygote.gradient(() -> loss(classifier, gen, tracked_steps), get_trainable_params(layers))
+    grads = Zygote.gradient(() -> loss(classifier, batches, tracked_steps), get_trainable_params(layers))
 
     println("Done grads")
     # discriminative step
@@ -179,9 +183,8 @@ end
 It contains main training loops for training a defined classifer for specified classes and data.
 Usage is discussed in the docs.
 """
-function train_classifier!(classifier::TextClassifier=TextClassifier(), data = (loadrecipe()["imdb"]))
+function train_classifier!(classifier::TextClassifier=TextClassifier(), batches=Nothing)
 
-    # dala_loader = imdb_classifier_data
     classes = 2
     hidden_layer_size = 50
     stlr_cut_frac=0.1
@@ -201,8 +204,8 @@ function train_classifier!(classifier::TextClassifier=TextClassifier(), data = (
 
     for epoch=1:epochs
         println("Epoch: $epoch")
-        gen = data
-        num_of_iters = numobs(data)
+        gen = batches
+        num_of_iters = length(batches)
         cut = num_of_iters * epochs * stlr_cut_frac
         for iter=1:num_of_iters
 
@@ -216,7 +219,7 @@ function train_classifier!(classifier::TextClassifier=TextClassifier(), data = (
             # Gradual-unfreezing Step with discriminative fine-tuning
             unfreezed_layers, cur_opts = (epoch < length(trainable)) ? (trainable[end-epoch+1:end], opts[end-epoch+1:end]) : (trainable, opts)
             println("start discriminative_step")
-            discriminative_step!(unfreezed_layers, classifier, gen, tracked_steps,ηL, cur_opts)
+            discriminative_step!(unfreezed_layers, classifier, batches, tracked_steps,ηL, cur_opts)
 
             println("End discriminative_step")