Struct GptJLMHeadModel

pub struct GptJLMHeadModel { /* private fields */ }

GPT-J Language Modeling head

GPT-J model with a decoding head (linear layer without bias). The weights of the linear layer are tied to the word embeddings It is made of the following blocks:

• transformer: Base GptJModel

Implementations

impl GptJLMHeadModel

pub fn new<'p, P>(p: P, config: &GptJConfig) -> GptJLMHeadModel

where P: Borrow<Path<'p>>,

Build a new GptJLMHeadModel

Arguments

• p - Variable store path for the root of the GPT-J model
• config - GptJConfig object defining the model architecture

Example

use rust_bert::gpt_j::{GptJConfig, GptJLMHeadModel};
use rust_bert::Config;
use std::path::Path;
use tch::{nn, Device};

let config_path = Path::new("path/to/config.json");
let device = Device::Cpu;
let p = nn::VarStore::new(device);
let config = GptJConfig::from_file(config_path);
let gpt_j: GptJLMHeadModel = GptJLMHeadModel::new(&p.root() / "gpt_j", &config);

pub fn forward_t( &self, input_ids: Option<&Tensor>, layer_past: Cache, attention_mask: Option<&Tensor>, token_type_ids: Option<&Tensor>, position_ids: Option<&Tensor>, input_embeds: Option<&Tensor>, _encoder_outputs: Option<&Tensor>, _decoder_input_ids: Option<&Tensor>, train: bool, ) -> Result<LMModelOutput, RustBertError>