coedit-base-webgpu-onnx

Flint CoEdIT Base for WebGPU (ONNX)

Browser-ready ONNX export of the 250M-parameter CoEdIT Base checkpoint, packaged for local sequence-to-sequence generation with Transformers.js and ONNX Runtime WebGPU.

Correct model lineage

This repository is an inference-format conversion, not a newly trained model and not an export of grammarly/coedit-large.

Grammarly publishes the CoEdIT dataset and the official Large, XL, and XXL checkpoints. It does not publish an official grammarly/coedit-base checkpoint. The Base checkpoint converted here is the FLAN-T5 Base fine-tune published by jbochi.

This repository contains only the two graphs required for cached browser generation:

• onnx/encoder_model.onnx
• onnx/decoder_model_merged.onnx

Tokenizer, model configuration, and generation configuration are included at the repository root. The redundant uncached decoder exports are intentionally omitted.

Usage with Transformers.js

import {
  AutoModelForSeq2SeqLM,
  AutoTokenizer,
} from "@huggingface/transformers";

const modelId = "imrahamed/coedit-base-webgpu-onnx";
const tokenizer = await AutoTokenizer.from_pretrained(modelId);
const model = await AutoModelForSeq2SeqLM.from_pretrained(modelId, {
  device: "webgpu",
  dtype: "fp32",
});

const input = await tokenizer(
  "Fix grammatical errors in this sentence: This are a test.",
);
const output = await model.generate({
  inputs: input.input_ids,
  attention_mask: input.attention_mask,
  max_new_tokens: 64,
});
console.log(
  tokenizer.decode(output.tolist()[0], {
    skip_special_tokens: true,
  }),
);

Expected output: This is a test.

CoEdIT task prompts

Export details

• Foundation architecture: FLAN-T5 Base
• Converted checkpoint: jbochi/coedit-base
• Fine-tuning dataset: grammarly/coedit
• Parameters: approximately 250M
• Format: ONNX, FP32
• Opset: 18
• Export task: text2text-generation-with-past
• Optimization: Optimum ONNX Runtime O2
• Intended execution provider: ONNX Runtime WebGPU
• CPU/WASM fallback should be provided by the consuming application.

The export was validated against the source model. Small floating-point differences from ONNX graph optimization may occur.

License and attribution

This derivative export follows the converted checkpoint's Apache 2.0 license. See the jbochi/coedit-base model card for its reported training details and metrics. CoEdIT paper and dataset attribution remains applicable.