# Training/Fine-tuning ## Data Preparation ### 1. Load the datasets Load `train.jsonl`, `validation.jsonl`, and `test.jsonl` splits. ```python raw_datasets = load_dataset("andreaceto/hasd") ``` --- ### 2. Create Label Mappings Now, we create the mappings from string labels (e.g., "schedule", "practitioner_name") to integer IDs. This is essential for training. We also need to create tags for the BIO (Beginning, Inside, Outside) entity scheme. ```python # --- Create Intent Label Mappings --- # Get all unique intent labels from the training data intent_labels = raw_datasets['train'].unique('intent') intent_labels.sort() # Sort for consistency id2intent = {i: label for i, label in enumerate(intent_labels)} intent2id = {label: i for i, label in enumerate(intent_labels)} print(f"Intent mapping (intent2id): {intent2id}\n") # --- Create Entity (NER) Label Mappings in BIO format --- # Get all unique entity labels entity_labels = ["appointment_id", "appointment_type", "practitioner_name"] # Create the full list of BIO tags ner_tags = ["O"] # 'O' for tokens outside any entity for label in entity_labels: ner_tags.append(f"B-{label}") # 'B' for Beginning of an entity ner_tags.append(f"I-{label}") # 'I' for Inside of an entity id2ner = {i: label for i, label in enumerate(ner_tags)} ner2id = {label: i for i, label in enumerate(ner_tags)} print(f"NER mapping (ner2id): {ner2id}") ``` --- ### 3. Preprocessing function This is the core function. It takes a single data example and does two things: 1. Tokenizes the text. 2. Aligns character-based entity spans (`start`, `end`) with the new wordpiece tokens, assigning the correct BIO tag ID to each token. ```python def preprocess_function(examples): # --- Intent Processing --- intent_ids = [intent2id[intent] for intent in examples['intent']] # --- Tokenization --- tokenized_inputs = tokenizer(examples['text'], truncation=True, is_split_into_words=False, return_offsets_mapping=True) # --- Entity (NER) Label Alignment --- ner_labels = [] for i, entities in enumerate(examples['entities']): word_ids = tokenized_inputs.word_ids(batch_index=i) label_ids = [ner2id["O"]] * len(word_ids) # For each entity, find the corresponding tokens and assign B- and I- tags for entity in entities: start_char, end_char, label = entity['start'], entity['end'], entity['label'] # This flag tracks if we've found the first token of the current entity first_token_of_entity_found = False for j, word_id in enumerate(word_ids): if word_id is None: continue token_char_span = tokenized_inputs['offset_mapping'][i][j] token_start, token_end = token_char_span # Check if the token is part of the entity if start_char < token_end and end_char > token_start: # This is the key change. We use the flag to decide the tag. if not first_token_of_entity_found: # This is the first token of the entity, assign the 'B-' tag label_ids[j] = ner2id[f"B-{label}"] first_token_of_entity_found = True else: # This is a subsequent token of the same entity, assign 'I-' label_ids[j] = ner2id[f"I-{label}"] ner_labels.append(label_ids) # Add the final processed labels to our tokenized inputs tokenized_inputs["intent_label"] = intent_ids tokenized_inputs["labels"] = ner_labels # Remove offset_mapping as it's not needed by the model tokenized_inputs.pop("offset_mapping", None) return tokenized_inputs ``` --- ### 4. Apply Preprocessing and Save Now we apply this function to our entire dataset and save the final, processed version. ```python # Apply the function to all splits of the dataset processed_datasets = raw_datasets.map(preprocess_function, batched=True, remove_columns=raw_datasets['train'].column_names) # Define the features for our processed dataset, including the new ClassLabels features = Features({ 'input_ids': Sequence(Value('int64')), 'attention_mask': Sequence(Value('int8')), 'intent_label': ClassLabel(names=list(intent2id.keys())), 'labels': Sequence(ClassLabel(names=list(ner2id.keys()))) }) # Cast the processed datasets to the defined features to include the label names processed_datasets = processed_datasets.cast(features) ``` --- ## Multitask Model ### 1. Multitask Model class To use the model you will need to define a `multitask_model.py` with the custom model class built upon our base model. ```python from transformers import AutoModel, PreTrainedModel import torch.nn as nn class MultitaskModel(PreTrainedModel): """ A custom Transformer model with two heads: one for intent classification and one for named entity recognition (token classification). """ config_class = AutoConfig def __init__(self, config): super().__init__(config) # Load the base transformer model (e.g., DistilBERT) self.transformer = AutoModel.from_config(config) # --- Heads --- # 1. Intent Classification Head (MLP) self.intent_classifier = nn.Sequential( nn.Linear(config.dim, config.dim // 2), nn.GELU(), # GELU is a smooth activation function, common in Transformers nn.Dropout(0.3), nn.Linear(config.dim // 2, self.config.num_intent_labels) ) # 2. NER (Token Classification) Head (MLP) self.ner_classifier = nn.Sequential( nn.Linear(config.dim, config.dim // 2), nn.GELU(), nn.Dropout(0.3), nn.Linear(config.dim // 2, self.config.num_ner_labels) ) # Dropout layer for regularization self.dropout = nn.Dropout(config.seq_classif_dropout) def forward( self, input_ids=None, attention_mask=None, intent_label=None, # For calculating intent loss labels=None, # For calculating NER loss ): # Get the last hidden states from the base transformer model outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask) sequence_output = outputs.last_hidden_state # Shape: (batch_size, sequence_length, hidden_size) # --- Intent Logits --- # Use the [CLS] token's output for intent classification cls_token_output = sequence_output[:, 0, :] cls_token_output = self.dropout(cls_token_output) intent_logits = self.intent_classifier(cls_token_output) # --- NER Logits --- # Use all token outputs for NER sequence_output = self.dropout(sequence_output) ner_logits = self.ner_classifier(sequence_output) # --- Calculate Combined Loss --- total_loss = 0 if intent_label is not None and labels is not None: loss_fct = nn.CrossEntropyLoss() # Intent loss intent_loss = loss_fct(intent_logits.view(-1, self.config.num_intent_labels), intent_label.view(-1)) # NER loss (ignore padding tokens with label -100) ner_loss = loss_fct(ner_logits.view(-1, self.config.num_ner_labels), labels.view(-1)) # Combine the losses (you can also weight them if one task is more important) total_loss = intent_loss + ner_loss return { "loss": total_loss, "intent_logits": intent_logits, "ner_logits": ner_logits, } ``` --- ### 2. Load Tokenizer ``` model_name = "distilbert-base-uncased" tokenizer = AutoTokenizer.from_pretrained(model_name) ``` --- ### 3. Custom Metrics Function This function is essential for a multitask model. It will be called by the Trainer at the end of each epoch to calculate both intent accuracy and NER F1-score. ```python def compute_metrics(eval_pred): # Unpack predictions and labels predictions, label_values = eval_pred intent_preds, ner_preds = predictions intent_labels, ner_labels = label_values # --- Intent Metrics --- intent_preds = np.argmax(intent_preds, axis=1) intent_accuracy = accuracy_score(intent_labels, intent_preds) intent_f1 = f1_score(intent_labels, intent_preds, average='weighted') # --- NER Metrics --- ner_preds = np.argmax(ner_preds, axis=2) # Remove padding tokens (where label is -100) and convert IDs to labels true_ner_labels = [] true_ner_predictions = [] id2ner = processed_datasets['train'].features['labels'].feature.names for i in range(len(ner_labels)): true_labels_row = [] true_predictions_row = [] for j in range(len(ner_labels[i])): if ner_labels[i][j] != -100: true_labels_row.append(id2ner[ner_labels[i][j]]) true_predictions_row.append(id2ner[ner_preds[i][j]]) true_ner_labels.append(true_labels_row) true_ner_predictions.append(true_predictions_row) ner_f1 = ner_f1_score(true_ner_labels, true_ner_predictions, mode='strict', scheme=IOB2) return { "intent_accuracy": intent_accuracy, "intent_f1": intent_f1, "ner_f1": ner_f1 } ``` --- ### 4. Instantiate the model We now create an instance of our `MultitaskModel`, passing it a configuration object that includes the number of labels for each head. ``` # --- Get DistilBert config file --- config = AutoConfig.from_pretrained(model_name) # --- Create Label Mappings Directly from Data --- # 1. Intent Labels # Get a sorted list of unique intent strings from the training set intent_label_list = processed_datasets['train'].features['intent_label'].names # Create the mappings id2intent = {i: label for i, label in enumerate(intent_label_list)} intent2id = {label: i for i, label in enumerate(intent_label_list)} # 2. NER Labels # Get a sorted list of unique entities strings from the training set ner_label_list = processed_datasets['train'].features['labels'].feature.names # Create the mappings id2ner = {i: label for i, label in enumerate(ner_label_list)} ner2id = {label: i for i, label in enumerate(ner_label_list)} # --- Add custom parameters to config object --- config.num_intent_labels = len(id2intent) config.num_ner_labels = len(id2ner) config.id2label_intent = id2intent config.label2id_intent = intent2id config.id2label_ner = id2ner config.label2id_ner = ner2id # --- Finally instantiate the model --- model = MultitaskModel(config) ``` --- # Inference ## Load the Model Load the trained tokenizer and model: ``` tokenizer = AutoTokenizer.from_pretrained("andreaceto/schedulebot-nlu-engine") multitask_model = MultitaskModel.from_pretrained("andreaceto/schedulebot-nlu-engine") ``` --- ## Preprocess raw text In order to being able to run inference on the model, it's necessary to use the tokenizer on the raw text in input: ``` inputs = self.tokenizer(text, return_tensors="pt") ``` --- ## Get predictions Now you can run inference using: ``` with torch.no_grad(): outputs = self.multitask_model(**inputs) ``` Since the model returns a dictionary, you can access logits by using: ``` intent_logits = outputs["intent_logits"] ner_logits = outputs["ner_logits"] ``` You can now extract prediction using `torch.argmax()` and converting the result into categorical by using the dictionaries in the config file (`id2label_intent` and `id2label_ner`). **N.B.**: Since SON format specifications require that all keys in an object be strings. When the transformers library saves your configuration, it correctly converts your integer keys (0, 1, 2, etc.) into strings ("0", "1", "2"), which causes a mismatch when trying to use extracted prediction with `torch.argmax()` as a key in the config dictionaries. The solution is to cast the `int` value of prediction as a string with `str()`.