@@ -10,6 +10,7 @@ import { createSession, type InferenceSession } from "../session/index.js"
1010import type { Tensor } from "../types.js"
1111import { verifyAndRead , parseManifest , type IMFManifest } from "./loader.js"
1212import { resolve } from "./registry.js"
13+ import { normalizeArabicInput , repetitionGuardCut } from "./guards.js"
1314import { EOS_ID , PAD_ID , decode , encode } from "./tokens.js"
1415
1516interface InputMeta {
@@ -22,6 +23,15 @@ interface MetadataSession extends InferenceSession {
2223 readonly inputMetadata ?: readonly InputMeta [ ]
2324}
2425
26+ export interface DecodeOptions {
27+ /** skip input normalization (raw text) */
28+ readonly raw ?: boolean
29+ /** streaming: called per emitted token (TODO.client-work 03) */
30+ readonly onToken ?: ( token : number , step : number ) => void
31+ /** per-step top1-top2 logit gap; low gap = low confidence (06) */
32+ readonly onConfidence ?: ( gap : number , step : number ) => void
33+ }
34+
2535export class IMFModel {
2636 readonly id : string
2737 private readonly manifest : IMFManifest
@@ -62,13 +72,15 @@ export class IMFModel {
6272 return IMFModel . fromZipBytes ( resolved . bytes )
6373 }
6474
65- async translate ( text : string , maxLen = 256 ) : Promise < string > {
66- const ids = encode ( text )
75+ async translate ( text : string , maxLen = 256 , opts : DecodeOptions = { } ) : Promise < string > {
76+ // models train on stripped input: normalize by default (TODO.client-work 02)
77+ const normalized = opts . raw === true ? text : normalizeArabicInput ( text )
78+ const ids = encode ( normalized )
6779 if ( ids . length === 1 ) return ""
6880 const hidden = await this . runEncoder ( ids )
6981 const tokens = this . kv
70- ? await this . greedyKv ( hidden , maxLen )
71- : await this . greedyPlain ( hidden , maxLen )
82+ ? await this . greedyKv ( hidden , maxLen , opts )
83+ : await this . greedyPlain ( hidden , maxLen , opts )
7284 return decode ( tokens )
7385 }
7486
@@ -123,25 +135,33 @@ export class IMFModel {
123135 return feeds
124136 }
125137
126- private argmaxLastStep ( logits : Tensor ) : number {
138+ private argmaxLastStep ( logits : Tensor ) : { token : number ; gap : number } {
127139 const dims = logits . dims
128140 const classes = dims [ dims . length - 1 ] !
129141 const data = logits . data as Float32Array | BigInt64Array
130142 const base = ( dims [ dims . length - 2 ] ! - 1 ) * classes
131143 let best = 0
132144 let bestVal = - Infinity
145+ let secondVal = - Infinity
133146 for ( let c = 0 ; c < classes ; c ++ ) {
134147 const v =
135148 typeof data [ base + c ] === "bigint" ? Number ( data [ base + c ] ) : ( data [ base + c ] as number )
136149 if ( v > bestVal ) {
150+ secondVal = bestVal
137151 bestVal = v
138152 best = c
153+ } else if ( v > secondVal ) {
154+ secondVal = v
139155 }
140156 }
141- return best
157+ return { token : best , gap : bestVal - secondVal }
142158 }
143159
144- private async greedyKv ( hidden : Tensor , maxLen : number ) : Promise < number [ ] > {
160+ private async greedyKv (
161+ hidden : Tensor ,
162+ maxLen : number ,
163+ opts : DecodeOptions = { } ,
164+ ) : Promise < number [ ] > {
145165 const generated : number [ ] = [ ]
146166 let current = [ PAD_ID ]
147167 let present : ReadonlyMap < string , Tensor > | undefined
@@ -161,9 +181,12 @@ export class IMFModel {
161181 } ,
162182 ...this . pastTensors ( present ) ,
163183 } )
164- const token = this . argmaxLastStep ( outputs [ "logits" ] ! )
184+ const { token, gap } = this . argmaxLastStep ( outputs [ "logits" ] ! )
165185 if ( token === EOS_ID ) break
166186 generated . push ( token )
187+ opts . onToken ?.( token , step )
188+ opts . onConfidence ?.( gap , step )
189+ if ( repetitionGuardCut ( generated , decode ( generated ) ) ) break
167190 present = new Map (
168191 this . pasts . map ( ( spec ) => [ spec . name , outputs [ spec . name . replace ( "past_" , "present_" ) ] ! ] ) ,
169192 )
@@ -172,7 +195,11 @@ export class IMFModel {
172195 return generated
173196 }
174197
175- private async greedyPlain ( hidden : Tensor , maxLen : number ) : Promise < number [ ] > {
198+ private async greedyPlain (
199+ hidden : Tensor ,
200+ maxLen : number ,
201+ opts : DecodeOptions = { } ,
202+ ) : Promise < number [ ] > {
176203 const generated : number [ ] = [ ]
177204 const decoderIds : number [ ] = [ PAD_ID ]
178205 for ( let step = 0 ; step < maxLen ; step ++ ) {
@@ -190,10 +217,13 @@ export class IMFModel {
190217 dims : hidden . dims ,
191218 } ,
192219 } )
193- const token = this . argmaxLastStep ( outputs [ "logits" ] ! )
220+ const { token, gap } = this . argmaxLastStep ( outputs [ "logits" ] ! )
194221 if ( token === EOS_ID ) break
195222 generated . push ( token )
196223 decoderIds . push ( token )
224+ opts . onToken ?.( token , step )
225+ opts . onConfidence ?.( gap , step )
226+ if ( repetitionGuardCut ( generated , decode ( generated ) ) ) break
197227 }
198228 return generated
199229 }
0 commit comments