Package ai.djl.modality.nlp.generate

Class CausalLMOutput

java.lang.Object

ai.djl.modality.nlp.generate.CausalLMOutput

public class CausalLMOutput
extends Object

CausalLMOuput is used to contain multiple output of a language model.

Constructor Summary

Constructors

Constructor	Description
CausalLMOutput(NDArray logits, NDArray hiddenState, NDList pastKeyValueList)	Constructs a new CausalLMOutput intance.
CausalLMOutput(NDArray logits, NDList pastKeyValues)	Constructs a new CausalLMOutput instance.

Method Summary

Modifier and Type	Method	Description
NDArray	getHiddenState()	Returns the value of the allHiddenStates.
NDArray	getLogits()	Returns the value of the logits.
NDList	getPastKeyValuesList()	Returns the value of the pastKeyValuesList.
void	setLogits(NDArray logits)	Sets the value of the logits.

Methods inherited from class java.lang.Object

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

Constructor Details

CausalLMOutput

public CausalLMOutput(NDArray logits,
 NDList pastKeyValues)

Constructs a new CausalLMOutput instance.

Parameters:

logits - the logits NDArray

pastKeyValues - the key-value cache

CausalLMOutput

public CausalLMOutput(NDArray logits,
 NDArray hiddenState,
 NDList pastKeyValueList)

Constructs a new CausalLMOutput intance.

Parameters:

logits - the logits NDArray

hiddenState - the first layer hiddenStates used as word embedding

pastKeyValueList - the key-value cache

Method Details

getLogits

public NDArray getLogits()

Returns the value of the logits.

Returns:

the value of logits

setLogits

public void setLogits(NDArray logits)

Sets the value of the logits.

Parameters:

logits - value of logits NDArray

getHiddenState

public NDArray getHiddenState()

Returns the value of the allHiddenStates.

Returns:

the value of allHiddenStates

getPastKeyValuesList

public NDList getPastKeyValuesList()

Returns the value of the pastKeyValuesList.

Returns:

the value of pastKeyValuesList