slightly change stream "abort" flow so token probabilities get successfully updated
| @@ -3194,7 +3194,8 @@ class StreamingProcessor { | |||
| 3194 | } | 3194 | } |
| 3195 | 3195 | ||
| 3196 | onStopStreaming() { | 3196 | onStopStreaming() { |
| 3197 | this.onErrorStreaming(); | 3197 | this.abortController.abort(); |
| 3198 | this.isFinished = true; | ||
| 3198 | } | 3199 | } |
| 3199 | 3200 | ||
| 3200 | /** | 3201 | /** |
| @@ -3239,9 +3240,13 @@ class StreamingProcessor { | |||
| 3239 | console.warn(`Stream stats: ${timestamps.length} tokens, ${seconds.toFixed(2)} seconds, rate: ${Number(timestamps.length / seconds).toFixed(2)} TPS`); | 3240 | console.warn(`Stream stats: ${timestamps.length} tokens, ${seconds.toFixed(2)} seconds, rate: ${Number(timestamps.length / seconds).toFixed(2)} TPS`); |
| 3240 | } | 3241 | } |
| 3241 | catch (err) { | 3242 | catch (err) { |
| 3243 | // in the case of a self-inflicted abort, we have already cleaned up | ||
| 3244 | if (!this.isFinished) | ||
| 3245 | { | ||
| 3242 | console.error(err); | 3246 | console.error(err); |
| 3243 | this.onErrorStreaming(); | 3247 | this.onErrorStreaming(); |
| 3244 | return; | 3248 | } |
| 3249 | return this.result; | ||
| 3245 | } | 3250 | } |
| 3246 | 3251 | ||
| 3247 | this.isFinished = true; | 3252 | this.isFinished = true; |
| @@ -4641,7 +4646,6 @@ export function stopGeneration() { | |||
| 4641 | let stopped = false; | 4646 | let stopped = false; |
| 4642 | if (streamingProcessor) { | 4647 | if (streamingProcessor) { |
| 4643 | streamingProcessor.onStopStreaming(); | 4648 | streamingProcessor.onStopStreaming(); |
| 4644 | streamingProcessor = null; | ||
| 4645 | stopped = true; | 4649 | stopped = true; |
| 4646 | } | 4650 | } |
| 4647 | if (abortController) { | 4651 | if (abortController) { |