Repository: xenova/transformers.js
File: src/models.js (JavaScript)

Code before:
// StableLm models
export class StableLMPreTrainedModel extends PreTrainedModel {
    /**
     * Creates a new instance of the `StableLMPreTrainedModel` class.
     * @param {Object} config The configuration of the model.
     * @param {any} session The ONNX session containing the model weights.
     * @param {GenerationConfig} generation_config The generation configuration.
     */
    constructor(config, session, generation_config) {
        super(config, session);
        this.generation_config = generation_config;

        // config doesn't contain pad_token_id, so we assume it is the eos_token_id
        this.config.pad_token_id = this.config.eos_token_id

        this.num_heads = this.config.num_attention_heads;
        this.num_layers = this.config.num_hidden_layers;
        this.dim_kv = this.config.hidden_size / this.num_heads;
    }
}

/**
 * The bare StableLM Model transformer outputting raw hidden-states without any specific head on top.
 */
export class StableLMModel extends StableLMPreTrainedModel { }

/**
 * StableLm Model with a `language modeling` head on top for Causal Language Modeling (with past).
 */
export class StableLmForCausalLM extends StableLmPreTrainedModel {
    /**
     * Calls the model on new inputs.
     * @param {Object} model_inputs The inputs to the model.
     * @returns {Promise<CausalLMOutput>} An object containing the model's output logits for causal language modeling.
     */
    async _call(model_inputs) {
        return new CausalLMOutputWithPast(await super._call(model_inputs));
    }
}
//////////////////////////////////////////////////

//////////////////////////////////////////////////
// AutoModels, used to simplify construction of PreTrainedModels
// (uses config to instantiate correct class)

/**
 * Base class of all AutoModels. Contains the `from_pretrained` function
 * which is used to instantiate pretrained models.
 */
export class PretrainedMixin {
    /**

Review comment (suggestion):
```suggestion
/**
 * The bare StableLm Model transformer outputting raw hidden-states without any specific head on top.
 */
export class StableLmModel extends StableLmPreTrainedModel { }

```

Code after:
// StableLm models
export class StableLmPreTrainedModel extends PreTrainedModel {
    /**
     * Creates a new instance of the `StableLmPreTrainedModel` class.
     * @param {Object} config The configuration of the model.
     * @param {any} session The ONNX session containing the model weights.
     * @param {GenerationConfig} generation_config The generation configuration.
     */
    constructor(config, session, generation_config) {
        super(config, session);
        this.generation_config = generation_config;

        // config doesn't contain pad_token_id, so we assume it is the eos_token_id
        this.config.pad_token_id = this.config.eos_token_id

        this.num_heads = this.config.num_attention_heads;
        this.num_layers = this.config.num_hidden_layers;
        this.dim_kv = this.config.hidden_size / this.num_heads;
    }
}

/**
 * The bare StableLm Model transformer outputting raw hidden-states without any specific head on top.
 */
export class StableLmModel extends StableLmPreTrainedModel { }

/**
 * StableLm Model with a `language modeling` head on top for Causal Language Modeling (with past).
 */
export class StableLmForCausalLM extends StableLmPreTrainedModel { }
//////////////////////////////////////////////////

//////////////////////////////////////////////////
// AutoModels, used to simplify construction of PreTrainedModels
// (uses config to instantiate correct class)

/**
 * Base class of all AutoModels. Contains the `from_pretrained` function
 * which is used to instantiate pretrained models.
 */
export class PretrainedMixin {
    /**
     * Mapping from model type to model class.
     * @type {Map<string, Object>[]}
     */
    static MODEL_CLASS_MAPPINGS = null;

    /**
     * Whether to attempt to instantiate the base class (`PretrainedModel`) if 
     * the model type is not found in the mapping.
     */