package kaun

  1. Overview
  2. Docs
Flax-inspired neural network library for OCaml

Install 

dune-project
 Dependency

github.com Readme Changelog Edit opam file Versions (2)

Authors

  1. Thibaut Mattio

Maintainers

  1. Thibaut Mattio

Sources

raven-1.0.0.alpha1.tbz
sha256=8e277ed56615d388bc69c4333e43d1acd112b5f2d5d352e2453aef223ff59867
sha512=369eda6df6b84b08f92c8957954d107058fb8d3d8374082e074b56f3a139351b3ae6e3a99f2d4a4a2930dd950fd609593467e502368a13ad6217b571382da28c

doc/kaun.models/Kaun_models/GPT2/For_causal_lm/index.html

Module GPT2.For_causal_lmSource

GPT-2 for causal language modeling

Sourceval create : ?config:config -> unit -> Kaun.module_
Sourceval forward : model:Kaun.module_ -> params:'a Kaun.params -> input_ids:(int32, Rune.int32_elt) Rune.t -> ?attention_mask:(int32, Rune.int32_elt) Rune.t -> ?position_ids:(int32, Rune.int32_elt) Rune.t -> ?labels:(int32, Rune.int32_elt) Rune.t -> training:bool -> unit -> (float, 'a) Rune.t * (float, 'a) Rune.t option

Returns (logits, loss) where logits has shape batch_size; seq_len; vocab_size