org.deeplearning4j.text.sentenceiterator.AggregatingSentenceIterator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of deeplearning4j-nlp Show documentation
There is a newer version: 1.0.0-M2.1
package org.deeplearning4j.text.sentenceiterator;

import lombok.NonNull;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * This is simple wrapper suited for aggregation of few SentenceIterators into single flow.
 *
 * @author [email protected]
 */
public class AggregatingSentenceIterator implements SentenceIterator {
    private List backendIterators;
    private SentencePreProcessor preProcessor;
    private AtomicInteger position = new AtomicInteger(0);

    private AggregatingSentenceIterator(@NonNull List list) {
        this.backendIterators = list;
    }

    @Override
    public String nextSentence() {
        if (!backendIterators.get(position.get()).hasNext() && position.get() < backendIterators.size()) {
            position.incrementAndGet();
        }

        return (preProcessor == null) ? backendIterators.get(position.get()).nextSentence() : preProcessor.preProcess(backendIterators.get(position.get()).nextSentence());
    }

    @Override
    public boolean hasNext() {
        for (SentenceIterator iterator: backendIterators) {
            if (iterator.hasNext()) {
                return true;
            }
        }
        return false;
    }

    @Override
    public void reset() {
        for (SentenceIterator iterator: backendIterators) {
            iterator.reset();
        }
        this.position.set(0);
    }

    @Override
    public void finish() {
        for (SentenceIterator iterator: backendIterators) {
            iterator.finish();
        }
    }

    @Override
    public SentencePreProcessor getPreProcessor() {
        return this.preProcessor;
    }

    @Override
    public void setPreProcessor(SentencePreProcessor preProcessor) {
        this.preProcessor = preProcessor;
    }

    public static class Builder {
        private List backendIterators = new ArrayList<>();
        private SentencePreProcessor preProcessor;

        public Builder () {

        }

        public Builder addSentenceIterator(@NonNull SentenceIterator iterator) {
            this.backendIterators.add(iterator);
            return this;
        }

        public Builder addSentenceIterators(@NonNull Collection iterator) {
            this.backendIterators.addAll(iterator);
            return this;
        }

        public Builder addSentencePreProcessor(@NonNull SentencePreProcessor preProcessor) {
            this.preProcessor = preProcessor;
            return this;
        }

        public AggregatingSentenceIterator build() {
            AggregatingSentenceIterator sentenceIterator = new AggregatingSentenceIterator(this.backendIterators);
            if (this.preProcessor != null) {
                sentenceIterator.setPreProcessor(this.preProcessor);
            }
            return sentenceIterator;
        }
    }
}