// ĠCopyright Ġ 2 0 0 0 - 2 0 2 0 ĠJetBrains Ġs . r . o . ĠUse Ġof Ġthis Ġsource Ġcode Ġis Ġgoverned Ġby Ġthe ĠApache Ġ 2 . 0 Ġlicense Ġthat Ġcan Ġbe Ġfound Ġin Ġthe ĠLICENSE Ġfile . Ċ package Ġcom . intellij . internal . ml . n gram Ċ Ċ import Ġcom . intellij . completion . n gram . s lp . modeling . n gram . NG ram Model Ċ import Ġcom . intellij . completion . n gram . s lp . trans lating . Vocab ulary Ċ import Ġcom . intellij . openapi . diagnostic . Logger Ċ import Ġorg . jetbrains . annotations . Test Only Ċ import Ġjava . io . External izable Ċ import Ġjava . io . IOException Ċ import Ġjava . io . Object Input Ċ import Ġjava . io . Object Output Ċ import Ġjava . util . concurrent . atomic . Atomic Integer Ċ import Ġkotlin . math . max Ċ Ċ class ĠV ocab ulary With Limit ( var Ġmax Vocab ulary Size : ĠInt , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġn Gram Order : ĠInt , Ġmax Sequence Size : ĠInt , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġsequence Initial Size : ĠInt Ġ= ĠSEQUENCE _ INITIAL _ SIZE ) Ġ: ĠV ocab ulary (), ĠExternal izable Ġ{ ĊĊĠ Ġcompanion Ġobject Ġ{ ĊĠĠĠ Ġinternal Ġval ĠLOG : ĠLogger Ġ= ĠLogger . getInstance ( Vocab ulary With Limit :: class . java ) ĊĊĠĠĠ Ġprivate Ġconst Ġval ĠSEQUENCE _ INITIAL _ SIZE : ĠInt Ġ= Ġ 1 0 0 ĊĠ Ġ} ĊĊĠ Ġprivate Ġval Ġcounter : ĠAtomicInteger Ġ= ĠAtomicInteger ( 1 ) ĊĊĠ Ġval Ġrecent : ĠNG ram Recent Tokens Ġ= ĠNG ram Recent Tokens ( max Sequence Size ) ĊĠ Ġval Ġrecent Sequence : ĠNG ram Recent Tokens Sequence Ġ= ĠNG ram Recent Tokens Sequence ( max Sequence Size , Ġn Gram Order , Ġsequence Initial Size ) ĊĊĠ Ġ/** ĊĠĠ Ġ* ĠIf Ġtoken Ġis Ġunknown , Ġit Ġwill Ġbe Ġtreated Ġas Ġ"< unknown Character >" ĊĠĠ Ġ*/ ĊĠ Ġfun Ġto Existing Indices ( token : ĠList < String >): ĠList < Int > Ġ{ ĊĠĠĠ Ġreturn Ġtoken . map NotNull Ġ{ Ġto Existing Index ( it ) Ġ} ĊĠ Ġ} ĊĊĠ Ġprivate Ġfun Ġto Existing Index ( token : ĠString ): ĠInt ? Ġ{ ĊĠĠĠ Ġreturn Ġword Indices [ token ] Ġ?: Ġword Indices [ unknown Character ] ĊĠ Ġ} ĊĊĠ Ġ/** ĊĠĠ Ġ* ĠIf Ġtoken Ġis Ġunknown , Ġvocabulary Ġwill Ġlearn Ġit . ĊĠĠ Ġ* ĊĠĠ Ġ* ĠIf Ġafter Ġlearning Ġnew Ġtoken Ġvocabulary Ġsize Ġexceeds ĠmaxSize Ġthreshold , Ġthe Ġoldest Ġtoken Ġwill Ġbe Ġforgot ten . ĊĠĠ Ġ*/ ĊĠ Ġfun Ġto Indices With Limit ( token : ĠList < String >, Ġmodel : ĠNG ram Model ): ĠList < Int > Ġ{ ĊĠĠĠ Ġval Ġindices Ġ= Ġtoken . map Ġ{ Ġto Index With Limit ( it , Ġmodel ) Ġ} ĊĊĠĠĠ Ġrecent Sequence . add With Limit ( model , Ġindices . last ()) ĊĠĠĠ Ġupdate Recent Tokens ( token ) ĊĠĠĠ Ġreturn Ġindices ĊĠ Ġ} ĊĊĠ Ġprivate Ġfun Ġto Index With Limit ( token : ĠString , Ġmodel : ĠNG ram Model ): ĠInt Ġ{ ĊĠĠĠ Ġvar Ġindex : ĠInt ? Ġ= Ġword Indices [ token ] ĊĠĠĠ Ġif Ġ( index Ġ== Ġnull ) Ġ{ ĊĠĠĠĠĠ Ġindex Ġ= Ġcounter . get And Increment () ĊĠĠĠĠĠ Ġword Indices [ token ] Ġ= Ġindex ĊĊĠĠĠĠĠ Ġif Ġ( recent . size () Ġ>= Ġmax Vocab ulary Size ) Ġ{ ĊĠĠĠĠĠĠĠ Ġval Ġ( to Remove , Ġlatest Appearance ) Ġ= Ġtrim Recent Tokens Size () ĊĠĠĠĠĠĠĠ Ġrecent Sequence . forget Until ( model , Ġif Ġ( recent . size () Ġ> Ġ 0 ) Ġrecent . last Index () Ġ- Ġlatest Appearance Ġelse Ġ 0 ) ĊĠĠĠĠĠĠĠ Ġfor Ġ( token ToRemove Ġin Ġto Remove ) Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġword Indices . remove ( token ToRemove ) ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠĠĠ Ġreturn Ġindex ĊĠ Ġ} ĊĊĠ Ġprivate Ġfun Ġtrim Recent Tokens Size (): ĠPair < ArrayList < String >, ĠInt > Ġ{ ĊĠĠĠ Ġval Ġto Remove : ĠArrayList < String > Ġ= Ġarray ListOf () ĊĠĠĠ Ġvar Ġlatest Appearance Ġ= Ġ 0 ĊĠĠĠ Ġwhile Ġ( recent . size () Ġ>= Ġmax Vocab ulary Size ) Ġ{ ĊĠĠĠĠĠ Ġval Ġ( token , Ġidx ) Ġ= Ġrecent . remove El dest () ĊĠĠĠĠĠ Ġlatest Appearance Ġ= Ġmax ( latest Appearance , Ġidx ) ĊĠĠĠĠĠ Ġto Remove . add ( token ) ĊĠĠĠ Ġ} ĊĠĠĠ Ġreturn Ġto Remove Ġto Ġlatest Appearance ĊĠ Ġ} ĊĊĠ Ġprivate Ġfun Ġupdate Recent Tokens ( tokens : ĠList < String >) Ġ{ ĊĠĠĠ Ġif Ġ( LOG . isDebugEnabled ) Ġ{ ĊĠĠĠĠĠ Ġassert Update Is Incremental ( tokens ) ĊĠĠĠ Ġ} ĊĊĠĠĠ Ġif Ġ( tokens . isNotEmpty ()) Ġ{ ĊĠĠĠĠĠ Ġrecent . update ( tokens . last ()) ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĊĠ Ġprivate Ġfun Ġassert Update Is Incremental ( tokens : ĠList < String >) Ġ{ ĊĠĠĠ Ġfor Ġ( i Ġin Ġ 0 .. tokens . size Ġ- Ġ 2 ) Ġ{ ĊĠĠĠĠĠ Ġassert ( recent . contains ( tokens [ i ])) Ġ{ ĊĠĠĠĠĠĠĠ Ġ" Cannot Ġfind Ġprevious Ġtoken Ġin Ġrecent : Ġ${ tokens [ i ]} Ġin Ġ$ tokens " ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĊĠ Ġ@ Throws ( IOException :: class ) ĊĠ Ġoverride Ġfun Ġwrite External ( out : ĠObject Output ) Ġ{ ĊĠĠĠ Ġout . writeInt ( max Vocab ulary Size ) ĊĠĠĠ Ġout . writeInt ( counter . get ()) ĊĠĠĠ Ġrecent . write External ( out ) ĊĠĠĠ Ġrecent Sequence . write External ( out ) ĊĊĠĠĠ Ġout . writeInt ( word Indices . size ) ĊĠĠĠ Ġfor Ġ(( token , Ġcode ) Ġin Ġword Indices ) Ġ{ ĊĠĠĠĠĠ Ġout . write Object ( token ) ĊĠĠĠĠĠ Ġout . writeInt ( code ) ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĊĠ Ġ@ Throws ( IOException :: class ) ĊĠ Ġoverride Ġfun Ġread External ( ins : ĠObject Input ) Ġ{ ĊĠĠĠ Ġmax Vocab ulary Size Ġ= Ġins . readInt () ĊĠĠĠ Ġcounter . set ( ins . readInt ()) ĊĠĠĠ Ġrecent . read External ( ins ) ĊĠĠĠ Ġrecent Sequence . read External ( ins ) ĊĊĠĠĠ Ġval Ġwords Size Ġ= Ġins . readInt () ĊĠĠĠ Ġfor Ġ( i Ġin Ġ 0 Ġuntil Ġwords Size ) Ġ{ ĊĠĠĠĠĠ Ġval Ġtoken Ġ= Ġins . read Object () Ġas ĠString ĊĠĠĠĠĠ Ġval Ġcode Ġ= Ġins . readInt () ĊĠĠĠĠĠ Ġword Indices [ token ] Ġ= Ġcode ĊĠĠĠ Ġ} ĊĠ Ġ} Ċ } Ċ Ċ /** Ċ Ġ* ĠStores Ġrecent Ġtokens Ġ( recent ) Ġwith Ġan Ġid Ġof Ġit Ġlast Ġappearance Ċ Ġ* Ċ Ġ* ĠLast Ġtoken Ġappearance Ġis Ġused Ġto Ġfind Ġa Ġminimum Ġsequence Ġwhich Ġhave Ġto Ġbe Ġforgot ten Ġtogether Ġwith Ġthe Ġtoken . Ċ Ġ*/ Ċ class ĠNG ram Recent Tokens ( private Ġval Ġmax Sequence Size : ĠInt ) Ġ: ĠExternal izable Ġ{ ĊĠ Ġprivate Ġvar Ġmax Token Index : ĠInt Ġ= ĠInt . MAX _ VALUE Ġ- Ġ 1 ĊĊĠ Ġprivate Ġval Ġnext Token Index : ĠAtomicInteger Ġ= ĠAtomicInteger ( 1 ) ĊĠ Ġprivate Ġval Ġrecent : ĠLinkedHashMap < String , ĠInt > Ġ= ĠLinkedHashMap () ĊĊĠ Ġ@ Synchron ized ĊĠ Ġfun Ġupdate ( token : ĠString ) Ġ{ ĊĠĠĠ Ġif Ġ( recent . containsKey ( token )) Ġ{ ĊĠĠĠĠĠ Ġrecent . remove ( token ) ĊĠĠĠ Ġ} ĊĠĠĠ Ġrecent [ token ] Ġ= Ġnext Token Index . get And Increment () ĊĊĠĠĠ Ġif Ġ( nextToken Index . get () Ġ> Ġmax Token Index ) Ġ{ ĊĠĠĠĠĠ Ġ// Ġshift Ġtoken Ġindex Ġto Ġavoid Ġtoken Ġindex Ġoverflow ĊĠĠĠĠĠ Ġreset Token Index () ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĊĠ Ġprivate Ġfun Ġreset Token Index () Ġ{ ĊĠĠĠ Ġval Ġfirst Ġ= Ġnext Token Index . get () Ġ- Ġmax Sequence Size ĊĠĠĠ Ġvar Ġnew Last Ġ= Ġ 0 ĊĠĠĠ Ġfor Ġ(( key , Ġvalue ) Ġin Ġrecent ) Ġ{ ĊĠĠĠĠĠ Ġval Ġnew Idx Ġ= Ġmax ( value Ġ- Ġfirst Ġ+ Ġ 1 , Ġ 0 ) ĊĠĠĠĠĠ Ġrecent [ key ] Ġ= Ġnew Idx ĊĠĠĠĠĠ Ġnew Last Ġ= Ġmax ( new Last , Ġnew Idx ) ĊĠĠĠ Ġ} ĊĠĠĠ Ġnext Token Index . set ( new Last Ġ+ Ġ 1 ) ĊĠ Ġ} ĊĊĠ Ġ@ Synchron ized ĊĠ Ġfun Ġremove El dest (): ĠPair < String , ĠInt > Ġ{ ĊĠĠĠ Ġval Ġ( token , Ġid ) Ġ= Ġrecent . iterator (). next () ĊĠĠĠ Ġrecent . remove ( token ) ĊĠĠĠ Ġreturn Ġtoken Ġto Ġid ĊĠ Ġ} ĊĊĠ Ġ@ Synchron ized ĊĠ Ġfun Ġcontains ( token : ĠString ): ĠBoolean Ġ= Ġrecent . contains ( token ) ĊĊĠ Ġ@ Synchron ized ĊĠ Ġfun Ġlast Index (): ĠInt Ġ= Ġnext Token Index . get () Ġ- Ġ 1 ĊĊĠ Ġ@ Synchron ized ĊĠ Ġfun Ġsize (): ĠInt Ġ= Ġrecent . size ĊĊĠ Ġ@ Synchron ized ĊĠ Ġ@ Throws ( IOException :: class ) ĊĠ Ġoverride Ġfun Ġwrite External ( out : ĠObject Output ) Ġ{ ĊĠĠĠ Ġout . writeInt ( nextToken Index . get ()) ĊĠĠĠ Ġout . writeInt ( recent . size ) ĊĠĠĠ Ġfor Ġ( entry Ġin Ġrecent ) Ġ{ ĊĠĠĠĠĠ Ġout . write Object ( entry . key ) ĊĠĠĠĠĠ Ġout . writeInt ( entry . value ) ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĊĠ Ġ@ Synchron ized ĊĠ Ġ@ Throws ( IOException :: class ) ĊĠ Ġoverride Ġfun Ġread External ( ins : ĠObject Input ) Ġ{ ĊĠĠĠ Ġnext Token Index . set ( ins . readInt ()) ĊĊĠĠĠ Ġval Ġrecent Size Ġ= Ġins . readInt () ĊĠĠĠ Ġfor Ġ( i Ġin Ġ 0 Ġuntil Ġrecent Size ) Ġ{ ĊĠĠĠĠĠ Ġrecent [ ins . read Object () Ġas ĠString ] Ġ= Ġins . readInt () ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĊĠ Ġ@ Test Only ĊĠ Ġfun Ġset Max Token Index ( new Max : ĠInt ) Ġ{ ĊĠĠĠ Ġmax Token Index Ġ= Ġnew Max ĊĠ Ġ} ĊĊĠ Ġ@ Test Only ĊĠ Ġfun Ġget Recent Tokens (): ĠList < Pair < String , ĠInt >> Ġ{ ĊĠĠĠ Ġval Ġrecent Tokens Ġ= Ġarray ListOf < Pair < String , ĠInt >>() ĊĠĠĠ Ġfor Ġ(( key , Ġvalue ) Ġin Ġrecent ) Ġ{ ĊĠĠĠĠĠ Ġrecent Tokens . add ( key Ġto Ġvalue ) ĊĠĠĠ Ġ} ĊĠĠĠ Ġreturn Ġrecent Tokens ĊĠ Ġ} Ċ }