/*
 *  * Copyright 2016 Skymind, Inc.
 *  *
 *  *    Licensed under the Apache License, Version 2.0 (the "License");
 *  *    you may not use this file except in compliance with the License.
 *  *    You may obtain a copy of the License at
 *  *
 *  *        http://www.apache.org/licenses/LICENSE-2.0
 *  *
 *  *    Unless required by applicable law or agreed to in writing, software
 *  *    distributed under the License is distributed on an "AS IS" BASIS,
 *  *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  *    See the License for the specific language governing permissions and
 *  *    limitations under the License.
 */

package org.datavec.nlp.vectorizer;


import org.datavec.api.berkeley.Counter;
import org.datavec.api.writable.IntWritable;
import org.datavec.api.records.reader.RecordReader;
import org.datavec.api.writable.Writable;
import org.datavec.nlp.reader.TfidfRecordReader;
import org.datavec.nlp.vectorizer.AbstractTfidfVectorizer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/**
 *
 * Nd4j tfidf vectorizer
 *
 * @author Adam Gibson
 */
public class TfidfVectorizer extends AbstractTfidfVectorizer<INDArray> {
    @Override
    public INDArray createVector(Object[] args) {
        INDArray ret = Nd4j.create(cache.vocabWords().size());
        Counter<String> docFrequencies = (Counter<String>)args[0];
        for(int i = 0; i < cache.vocabWords().size(); i++) {
            double freq = docFrequencies.getCount(cache.wordAt(i));
            double tfidf = cache.tfidf(cache.wordAt(i),freq);
            ret.putScalar(i,tfidf);
        }
        return ret;
    }

    @Override
    public INDArray fitTransform(RecordReader reader) {
        return fitTransform(reader,null);
    }

    @Override
    public INDArray fitTransform(final RecordReader reader, RecordCallBack callBack) {
        final List<Collection<Writable>> records = new ArrayList<>();
        final TfidfRecordReader reader2 = (TfidfRecordReader) reader;
        fit(reader,new RecordCallBack() {
            @Override
            public void onRecord(Collection<Writable> record) {
                if(reader.getConf().get(TfidfRecordReader.APPEND_LABEL).equals("true")) {
                    record.add(new IntWritable(reader2.getCurrentLabel()));
                }
                records.add(record);
            }
        });

        if(records.isEmpty())
            throw new IllegalStateException("No records found!");
        INDArray ret = Nd4j.create(records.size(),cache.vocabWords().size());
        int i = 0;
        for(Collection<Writable> record : records) {
            ret.putRow(i++, transform(record));
            if(callBack != null) {
                callBack.onRecord(record);
            }
        }

        return ret;
    }

    @Override
    public INDArray transform(Collection<Writable> record) {
        Counter<String> wordFrequencies = wordFrequenciesForRecord(record);
        return createVector(new Object[]{wordFrequencies});

    }
}

@@ -763,80 +763,80 @@
 api.
-writable.IntWritable;%0Aimport org.datavec.api.records.reader.RecordReader
+records.reader.RecordReader;%0Aimport org.datavec.api.writable.IntWritable
 ;%0Aim
@@ -928,67 +928,8 @@
 er;%0A
-import org.datavec.nlp.vectorizer.AbstractTfidfVectorizer;%0A
 impo
@@ -1291,71 +1291,8 @@
 ) %7B%0A
-        INDArray ret = Nd4j.create(cache.vocabWords().size());%0A
     
@@ -1354,16 +1354,81 @@
 rgs%5B0%5D;%0A
+        double%5B%5D vector = new double%5Bcache.vocabWords().size()%5D;%0A
         
@@ -1564,20 +1564,17 @@
     
-double tfidf
+vector%5Bi%5D
  = c
@@ -1619,64 +1619,44 @@
     
-    ret.putScalar(i,tfidf);%0A        %7D%0A        return ret
+%7D%0A        return Nd4j.create(vector)
 ;%0A  

Use java array for element-wise creation