change input to vector of ngrams, update tests

Author: pazzo83

Project.toml

@@ -20,4 +20,4 @@ MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["MLJBase", "Test"]
+test = ["MLJBase", "Test", "TextAnalysis"]

src/MLJText.jl

@@ -40,26 +40,13 @@ MMI.@mlj_model mutable struct TfidfTransformer <: MLJModelInterface.Unsupervised
     max_doc_freq::Float64 = 0.98
     min_doc_freq::Float64 = 0.02
     smooth_idf::Bool = true
-    min_ngram_range::Int = 1
-    max_ngram_range::Int = 1
 end
 
 struct TfidfTransformerResult
     vocab::Vector{String}
     idf_vector::Vector{Float64}
 end
 
-_build_corpus(transformer::TfidfTransformer, docs::Vector{String}) = _build_corpus(transformer, StringDocument.(docs))
-
-function _build_corpus(transformer::TfidfTransformer, docs::Vector{StringDocument{String}})
-    corpus = Corpus(
-        NGramDocument.(
-            ngrams.(docs, transformer.min_ngram_range, transformer.max_ngram_range)
-        )
-    )
-    return corpus
-end
-
 function limit_features(doc_term_matrix::DocumentTermMatrix, high::Int, low::Int)
     doc_freqs = vec(sum(doc_term_matrix.dtm, dims=1))
 
@@ -77,7 +64,7 @@ function limit_features(doc_term_matrix::DocumentTermMatrix, high::Int, low::Int
     return (doc_term_matrix.dtm[:, mask], new_terms)
 end
 
-MMI.fit(transformer::TfidfTransformer, verbosity::Int, X) = _fit(transformer, verbosity, _build_corpus(transformer, X))
+MMI.fit(transformer::TfidfTransformer, verbosity::Int, X) = _fit(transformer, verbosity, Corpus(NGramDocument.(X)))
 
 function _fit(transformer::TfidfTransformer, verbosity::Int, X::Corpus)
     transformer.max_doc_freq < transformer.min_doc_freq && error("Max doc frequency cannot be less than Min doc frequency!")
@@ -131,7 +118,7 @@ function build_tfidf!(dtm::SparseMatrixCSC{T}, tfidf::SparseMatrixCSC{F}, idf_ve
     return tfidf
 end
 
-MMI.transform(transformer::TfidfTransformer, result::TfidfTransformerResult, v) = _transform(transformer, result, _build_corpus(transformer, v))
+MMI.transform(transformer::TfidfTransformer, result::TfidfTransformerResult, v) = _transform(transformer, result, Corpus(NGramDocument.(v)))
 
 function _transform(::TfidfTransformer, result::TfidfTransformerResult, v::Corpus)
     m = DocumentTermMatrix(v, result.vocab)
@@ -161,7 +148,7 @@ MMI.metadata_pkg(TfidfTransformer,
 )
 
 MMI.metadata_model(TfidfTransformer,
-               input_scitype = AbstractVector{STB.Textual},
+               input_scitype = AbstractVector{STB.Multiset{STB.Textual}},
                output_scitype = AbstractMatrix{STB.Continuous},# ie, a classifier
                docstring = "Build TF-IDF matrix from raw documents",         # brief description
                path = "MLJText.TfidfTransformer"

test/runtests.jl

@@ -1,24 +1,35 @@
 using MLJText # substitute for correct interface pkg name
 using Test
 using MLJBase
+using TextAnalysis
 
 @testset "tfidf transformer" begin
     # add some test docs
     docs = ["Hi my name is Sam.", "How are you today?"]
 
+    # convert to ngrams
+    ngram_vec = ngrams.(documents(Corpus(NGramDocument.(docs))))
+
+    # train transformer
     tfidf_transformer = MLJText.TfidfTransformer()
-    test = machine(tfidf_transformer, docs)
-    fit!(test)
+    test = machine(tfidf_transformer, ngram_vec)
+    MLJ.fit!(test)
 
-    test1 = transform(test, ["Another sentence ok"])
+    # test
+    test_doc = ngrams(NGramDocument("Another sentence ok"))
+    transform(test, [test_doc])
     @test sum(test1, dims=2)[1] == 0.0
     @test size(test1) == (1, 11)
 
-    test2 = transform(test, ["Listen Sam, today is not the day."])
+    test_doc2 = ngrams(NGramDocument("Listen Sam, today is not the day."))
+    transform(test, [test_doc2])
     @test sum(test2, dims=2)[1] > 0.0
     @test size(test2) == (1, 11)
 
-    test3 = transform(test, ["Another sentence ok", "Listen Sam, today is not the day."])
+    test_doc3 = ngrams.(
+        Corpus(NGramDocument("Another sentence ok"), NGramDocument("Listen Sam, today is not the day."))
+    )
+    transform(test, test_doc3)
     @test sum(test3, dims=2)[1] == 0.0
     @test sum(test3, dims=2)[2] > 0.0
     @test size(test3) == (2, 11)