Convert model.safetensors in order to be able to execute it with ExecuteTorch: how to prepare the example input and dynamic shape information?

Hi!

I've trained for fine-tuning the Bert model to use it for Named Entity Recognition.

Now I want to convert the resulting model.safetensors in order to be able to execute it with ExecuteTorch. Thanks to the explanation of a kind guy : https://dev-discuss.pytorch.org/t/what-is-the-correct-future-proof-way-of-deploying-a-pytorch-python-model-in-c-for-inference/2775/11?u=raphael10-collab ,
I've learned that, in order to export the torch.nn.Module into aExportedProgram, I need first to prepare the example input and dynamic shape information.

So.... my question is: which dynamic shape information should I use, since the model.safetensors I produced is just a fine-tuning of the Bert Model?
Should I use the shapes from here: https://github.com/google-research/bert/blob/master/modeling.py#L389 : input_ids: int32 Tensor of shape [batch_size, seq_length] containing word ids ?

This the code I used to fine-tune Bert model for NER task:

BERT-NER.py :

# https://github.com/tozameerkhan/Fine-Tuning-BERT-for-Named-Entity-Recognition/blob/main/BERTfineTunningFinal.ipynb

# 1. Setup and Installation

import datasets
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from transformers import BertTokenizerFast
from transformers import DataCollatorForTokenClassification
from transformers import TrainingArguments, Trainer, EarlyStoppingCallback
from transformers import logging as hf_logging
from transformers import pipeline
import json
from pprint import pprint
from torchmetrics.text.bert import BERTScore


bertscore = BERTScore()

hf_logging.set_verbosity_info() #to display informational messages.

from transformers import AutoModelForTokenClassification

import warnings
warnings.filterwarnings('ignore')

import matplotlib.pyplot as plt
plt.style.use("fivethirtyeight")


# 2. Data Exploration (EDA)

# Load Dataset

conll2003 = datasets.load_dataset("conll2003", trust_remote_code=True)
conll2003

# Convert to DataFrame
train_df = pd.DataFrame(conll2003['train'])
validation_df = pd.DataFrame(conll2003['validation'])
test_df = pd.DataFrame(conll2003['test'])

# Data Overview

print(train_df.head())
print(f"Number of sentences in the training set: {len(train_df)}")
print(f"Number of sentences in the validation set: {len(validation_df)}")
print(f"Number of sentences in the test set: {len(test_df)}")

label_list = conll2003["train"].features["ner_tags"].feature.names
print(label_list)

# Distribution of Sentence Lengths
train_df['sentence_length'] = train_df['tokens'].apply(len)
plt.figure(figsize=(10, 6))
sns.histplot(train_df['sentence_length'], bins=30, kde=True)
plt.title('Distribution of Sentence Lengths in Training Set')
plt.xlabel('Sentence Length')
plt.ylabel('Frequency')
plt.show()

# Distribution of Named Entity Tags
ner_tags = conll2003['train'].features['ner_tags'].feature.names
tag_counts = [0] * len(ner_tags)
for tags in train_df['ner_tags']:
    for tag in tags:
        tag_counts[tag] += 1

plt.figure(figsize=(12, 6))
sns.barplot(x=ner_tags, y=tag_counts)
plt.title('Distribution of Named Entity Tags in Training Set')
plt.xlabel('Named Entity Tag')
plt.ylabel('Count')
plt.xticks(rotation=45)
plt.show()

# 3. Data Preparation

# Tokenization and Label Alignment

#load a pre-trained tokenizer.
tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")

example_1 = conll2003['train'][0]
tokenized_input = tokenizer(example_1["tokens"], is_split_into_words=True)
tokens = tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])
word_ids = tokenized_input.word_ids()
print("word_ids :: ",word_ids)
''' As we can see, it returns a list with the same number of elements as our processed input ids, 
    mapping special tokens to None and all other tokens to their respective word.'''
print()#Function to tokenize and align labels with respect to the tokens.
def tokenize_and_align_labels(examples, label_all_tokens=True):
    tokenized_inputs = tokenizer(examples['tokens'], truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples['ner_tags']):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(label[word_idx] if label_all_tokens else -100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_datasets = conll2003.map(tokenize_and_align_labels, batched=True)

q = tokenize_and_align_labels(conll2003['train'][3:4])
print(q)

for token, label in zip(tokenizer.convert_ids_to_tokens(q["input_ids"][0]),q["labels"][0]):
    print(f"{token:_<40} {label}")


# 4. Model Fine-Tuning

try:
    model = AutoModelForTokenClassification.from_pretrained("bert-base-uncased", num_labels=len(conll2003['train'].features['ner_tags'].feature.names))
    print("Model loaded successfully")
except Exception as e:
    # Print the exception if there is an error
    print(f"Error loading model: {e}")

data_collator = DataCollatorForTokenClassification(tokenizer)
#Batch Creation, Dynamic Padding, Attention Masks

#Define early stopping callback
early_stopping = EarlyStoppingCallback(
    early_stopping_patience = 2  #Number of epochs to wait for improvement
)

# Define Training Arguments
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=10,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    load_best_model_at_end=True,
    save_total_limit=3,
)

# Initialize Trainer with EarlyStoppingCallback

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
    tokenizer=tokenizer,
    data_collator=data_collator,
    callbacks=[early_stopping]
)

train_result = trainer.train()


model.save_pretrained("ner_model")
tokenizer.save_pretrained("tokenizer")


id2label = {
    str(i): label for i,label in enumerate(label_list)
}

label2id = {
    label: str(i) for i,label in enumerate(label_list)
}

config = json.load(open("ner_model/config.json"))
config["id2label"] = id2label
config["label2id"] = label2id
json.dump(config, open("ner_model/config.json","w"))
# 5. Visualization of Training Process

# Training and Evaluation Loss Curves

eval_results = trainer.evaluate()
print(f"Evaluation results :: {eval_results}")

predictions, labels, _ = trainer.predict(tokenized_datasets["test"])
predictions = np.argmax(predictions, axis=2)

# Remove ignored index (special tokens)
true_predictions = [
            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
]

true_labels = [
            [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
]

#pprint(bertscore(true_predictions, true_labels))


train_losses = []

# Extract training losses from log history
for log in trainer.state.log_history:
    if 'loss' in log:
        train_losses.append(log['loss'])

# Plot training losses
plt.figure(figsize=(10, 6))
plt.plot(train_losses, label='Training Loss')
plt.xlabel('Steps')
plt.ylabel('Loss')
plt.title('Training Loss Curves')
plt.legend()
plt.show()



eval_losses = [] 

for log in trainer.state.log_history:
    # Check if the log contains validation loss information
    if 'eval_loss' in log:
        # Append the validation loss to the list
        eval_losses.append(log['eval_loss'])        


# Plot Validation losses
plt.figure(figsize=(10, 6))
plt.plot(eval_losses[0:6], label='Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Validation Loss Curves')
plt.legend()
plt.xticks(range(1, 6))
plt.show()


# 6. Test the Model with Example Texts

from transformers import AutoModelForTokenClassification
from transformers import pipeline #provides an easy-to-use interface for performing NLP tasks using pre-trained models
from transformers import BertTokenizerFast

model= AutoModelForTokenClassification.from_pretrained("ner_model")
tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")

nlp = pipeline("ner", model=model, tokenizer=tokenizer)

text = "Steve Jobs  is the Founder of Apple."
ner_results = nlp(text)
print(ner_results)

text = "Shortage of Jobs in India."
ner_results = nlp(text)
print(ner_results)

text = "London, the capital of England and the United Kingdom, is a 21st-century city with history \
         stretching back to Roman times"
ner_results = nlp(text)
print(ner_results)

text = "Delhi, India’s capital territory, is a massive metropolitan area in the country’s north. \
        In Old Delhi, a neighborhood dating to the 1600s, stands the imposing Red Fort."
ner_results = nlp(text)
print(ner_results)


# https://pytorch.org/docs/stable/export.html
import torch
from torch.export import export


print("tokenized_input :: ",tokenized_input)

len(example_1['ner_tags']), len(tokenized_input["input_ids"])

And this is the code I drafted to convert the resulting model.safetensors in order to be able to execute it with ExecuteTorch:

loadSafetensorIntoBERTModel.py :

# https://dev-discuss.pytorch.org/t/what-is-the-correct-future-proof-way-of-deploying-a-pytorch-python-model-in-c-for-inference/2775/11?u=raphael10-collab

import torch
from safetensors.torch import load_model
from torch.export import Dim, export

load_model(model, "./ner_model/model.safetensors")

# prepare the example input and dynamic shape information


# https://github.com/google-research/bert
# https://arxiv.org/abs/1810.04805
# https://arxiv.org/pdf/1810.04805
# https://github.com/tozameerkhan/Fine-Tuning-BERT-for-Named-Entity-Recognition/blob/main/BERTfineTunningFinal.ipynb
# https://huggingface.co/docs/transformers/model_doc/bert
# https://huggingface.co/google-bert/bert-base-cased

#example_args = torch.randn(batch_size, seq_length)
#ep = torch.export.export(model, example_args);

example_args = (torch.randn(batch_size, seq_length), )

# Create a dynamic batch size
batch = Dim("batch")


exported_program: torch.export.ExportedProgram = export(
    Mod(), args=example_args
)
print(exported_program)

What am I missing and/or doing wrong?
How to make it work?

cc @mergennachin @byjlw