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

Code before:
// Segformer models
export class SegformerPreTrainedModel extends PreTrainedModel { }

/**
 * The bare SegFormer encoder (Mix-Transformer) outputting raw hidden-states without any specific head on top.
 */
export class SegformerModel extends SegformerPreTrainedModel { }

/**
 * SegFormer Model transformer with an image classification head on top (a linear layer on top of the final hidden states) e.g. for ImageNet.
 */
export class SegformerForImageClassification extends SegformerPreTrainedModel { }

/**
 * SegFormer Model transformer with an all-MLP decode head on top e.g. for ADE20k, CityScapes.
 */
export class SegformerForSemanticSegmentation extends SegformerPreTrainedModel { }

//////////////////////////////////////////////////

//////////////////////////////////////////////////
// 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 {

Review comment (suggestion):
```suggestion
// StableLM models
export class StableLmPreTrainedModel extends PreTrainedModel {
    /**
     * Creates a new instance of the `StableLmPreTrainedModel` class.
     * @param {Object} config The configuration of the model.
```

Code after:
// Segformer models
export class SegformerPreTrainedModel extends PreTrainedModel { }

/**
 * The bare SegFormer encoder (Mix-Transformer) outputting raw hidden-states without any specific head on top.
 */
export class SegformerModel extends SegformerPreTrainedModel { }

/**
 * SegFormer Model transformer with an image classification head on top (a linear layer on top of the final hidden states) e.g. for ImageNet.
 */
export class SegformerForImageClassification extends SegformerPreTrainedModel { }

/**
 * SegFormer Model transformer with an all-MLP decode head on top e.g. for ADE20k, CityScapes.
 */
export class SegformerForSemanticSegmentation extends SegformerPreTrainedModel { }

//////////////////////////////////////////////////

//////////////////////////////////////////////////
// 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 { }