Package ai.djl.modality.nlp.generate

Class ContrastiveSeqBatchScheduler

java.lang.Object

ai.djl.modality.nlp.generate.SeqBatchScheduler

ai.djl.modality.nlp.generate.ContrastiveSeqBatchScheduler

public class ContrastiveSeqBatchScheduler
extends SeqBatchScheduler

ContrastiveSeqBatchScheduler is a class which implements the contrastive search algorithm used in SeqBatchScheduler.

Constructor Summary

Constructors

Constructor	Description
ContrastiveSeqBatchScheduler(Predictor<NDList,CausalLMOutput> lmBlock, SearchConfig config)	Constructs a new ContrastiveSeqBatchScheduler instance.

Method Summary

Modifier and Type	Method	Description
NDArray	inferenceCall()	An inference call in an iteration.
SeqBatcher	initForward(NDArray inputIds, NDArray batchUids)	Initializes the iteration and SeqBatcher.

Methods inherited from class ai.djl.modality.nlp.generate.SeqBatchScheduler

addRequest, collectResults, incrementForward

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

ContrastiveSeqBatchScheduler

public ContrastiveSeqBatchScheduler(Predictor<NDList,CausalLMOutput> lmBlock,
 SearchConfig config)

Constructs a new ContrastiveSeqBatchScheduler instance.

Parameters:

lmBlock - the predictor containing language model

config - the autoregressive search configuration

Method Details

initForward

public SeqBatcher initForward(NDArray inputIds,
 NDArray batchUids)
                       throws TranslateException

Initializes the iteration and SeqBatcher.

Specified by:

initForward in class SeqBatchScheduler

Parameters:

inputIds - the input token ids.

batchUids - the request uid identifying a sequence

Returns:

SeqBatcher Stores the search state and operate on the BatchTensorList

Throws:

TranslateException - if forward fails

inferenceCall

public NDArray inferenceCall()
                      throws TranslateException

An inference call in an iteration.

Specified by:

inferenceCall in class SeqBatchScheduler

Returns:

the output token ids

Throws:

TranslateException - if forward fails