Allow negative indices for slicing

Author: xenova

src/models.js

@@ -64,7 +64,6 @@ import {
 
 import {
     Tensor,
-    cat
 } from './utils/tensor.js';
 
 import { executionProviders, ONNX } from './backends/onnx.js';
@@ -839,16 +838,9 @@ export class PreTrainedModel extends Callable {
                 // In most cases, this will be [batch_size, 1, vocab_size]
                 // So, we select the last token's logits:
                 // (equivalent to `logits = outputs.logits[:, -1, :]`)
-                let extractedLogits = [];
-                for (const batch of output.logits) {
-                    // Extract logits corresponding to the last token
-                    let lastLogits = batch[-1];
-
-                    // Add back batch dimension (needed for `cat`)
-                    lastLogits.dims = [1, ...lastLogits.dims];
-                    extractedLogits.push(lastLogits)
-                }
-                let logits = cat(extractedLogits);
+                let logits = output.logits.slice(null, -1, null);
+
+                // Apply logits processor
                 logits_processor(beam.output_token_ids, logits);
 
                 let sampledTokens = sampler(logits);

src/utils/tensor.js

@@ -185,9 +185,12 @@ export class Tensor extends ONNXTensor {
                 newTensorDims.push(this.dims[sliceIndex]);
 
             } else if (typeof slice === 'number') {
-                if (slice < 0 || slice >= this.dims[sliceIndex]) {
+                if (slice < -this.dims[sliceIndex] || slice >= this.dims[sliceIndex]) {
                     throw new Error(`IndexError: index ${slice} is out of bounds for dimension ${sliceIndex} with size ${this.dims[sliceIndex]}`);
                 }
+                if (slice < 0) {
+                    slice += this.dims[sliceIndex];
+                }
 
                 // A number means take a single element
                 newOffsets.push([slice, slice + 1]);