Add StableLM support

[D4ve-R]

Adds support for

• StableLMForCausalLM
• StableLMForSequenceClassification

Closes #549

[D4ve-R]

StableLM support in Optimum is currently underway.
huggingface/optimum#1719

[xenova]

I know it's still a draft, but I added a commit to get it into a working state :)

Example usage:

import { pipeline } from '@xenova/transformers';

const generator = await pipeline('text-generation', 'Xenova/tiny-random-StableLmForCausalLM')
const output = await generator('hi')
console.log(output);

It's just a randomly initialized tiny model, so the output is gibberish. I will export some larger models too.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[D4ve-R]

Awesome! 😎
@xenova, how did you create the testnet? Just export a random pytorch net with the same architecture as stablelm?

[xenova]

Yes that is how you do it - but luckily in this case, we already have a tiny stablelm model: https://huggingface.co/hf-internal-testing/tiny-random-StableLmForCausalLM (check out the org for the full list of them)

[xenova]

Here are some of the larger models:

• https://huggingface.co/Xenova/stablelm-2-1_6b
• https://huggingface.co/Xenova/stablelm-2-zephyr-1_6b

which we can test with

(see all)

[xenova]

Seems to work alright! 👍

import { pipeline } from '@xenova/transformers';

const generator = await pipeline('text-generation', 'Xenova/stablelm-2-zephyr-1_6b')

const prompt = [{'role': 'user', 'content': 'Tell me a joke'}];
const inputs = generator.tokenizer.apply_chat_template(prompt, { add_generation_prompt: true, tokenize: false });

const output = await generator(inputs, { max_new_tokens: 20 })
console.log(output[0].generated_text);
// "<|user|>\nTell me a joke\n<|assistant|>\nWhy did the tomato turn red?\n\nBecause it saw the salad dressing!"

[D4ve-R]

Really cool! Will turn this into a pr.

[xenova]

Thanks! I think the last thing to do is just export and test with some sequence classifier models. Is that something you'd like to work on?

[D4ve-R]

Yes I will do that. Just to check, what do you mean by export?

[xenova]

what do you mean by export?

export / convert to onnx 👍

[xenova]

Looks like there aren't any stablelm text-classification models on the HF Hub (other than your test of course). So I think it will be a good idea to move that to a separate PR and get this merged?

[xenova]

TODO: move stablelm text classification to separate PR

[D4ve-R]

Wow, seems like you can read my mind 😆
You're right, will move this to another PR!
My response i had typed in, but not send yet:

Little update, since there is currently no stablelm model for text classification on the hub, I tried training my own.
Unfortunately I was unable to train on 16GB T4. I will add the code below to train stablelm-2-1_6b on twitter-financial-news-sentiment, if somebody wants to give it a shot.
I think we should split this in two seperate prs and merged the Seq Classification one when a model is available.

import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification, TrainingArguments, Trainer
import evaluate
from huggingface_hub import login
login()

dataset = "zeroshot/twitter-financial-news-sentiment"
model = "stabilityai/stablelm-2-1_6b"

dataset = load_dataset(dataset)
tokenizer = AutoTokenizer.from_pretrained(model)
tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)

tokenized_dataset = dataset.map(preprocess_function, batched=True)
id2label = {0: "Bearish", 1: "Bullish", 2: "Neutral"}
label2id = {"Bearish": 0, "Bullish": 1, "Neutral": 2}

model = AutoModelForSequenceClassification.from_pretrained(
    model,
    num_labels=len(id2label),
    id2label=id2label,
    label2id=label2id,
)
model.config.pad_token_id = model.config.eos_token_id

training_args = TrainingArguments(
    output_dir="stablelm-2-1_6b-sentiment",
    learning_rate=2e-5,
    # TODO: will it fit 16GB?
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=2,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    # TODO: will it fit 16GB?
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    optim="adamw_bnb_8bit",
    load_best_model_at_end=True,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()
trainer.push_to_hub()

[xenova]

naming nits

[xenova]

Merged! Thanks for this @D4ve-R! 🤗

Example: Text generation with Xenova/stablelm-2-zephyr-1_6b.

import { pipeline } from '@xenova/transformers';

// Create text generation pipeline
const generator = await pipeline('text-generation', 'Xenova/stablelm-2-zephyr-1_6b');

// Define the prompt and list of messages
const prompt = "Tell me a funny joke."
const messages = [
    { "role": "system", "content": "You are a helpful assistant." },
    { "role": "user", "content": prompt },
]

// Apply chat template
const inputs = generator.tokenizer.apply_chat_template(messages, {
    tokenize: false,
    add_generation_prompt: true,
});

// Generate text
const output = await generator(inputs, { max_new_tokens: 20 });
console.log(output[0].generated_text);
// "<|system|>\nYou are a helpful assistant.\n<|user|>\nTell me a funny joke.\n<|assistant|>\nHere's a joke for you:\n\nWhy don't scientists trust atoms?\n\nBecause they make up everything!"