Update how_to_use.md

how_to_use.md

@@ -1,4 +1,4 @@
-# Intended Use
+# Training/Fine-tuning
 
 ## Data Preparation
 
@@ -48,43 +48,40 @@ This is the core function. It takes a single data example and does two things:
 ```python
 def preprocess_function(examples):
     # --- Intent Processing ---
-    # Convert intent strings to integer IDs
     intent_ids = [intent2id[intent] for intent in examples['intent']]
 
     # --- Tokenization ---
-    # Tokenize the text. `truncation=True` and `padding` are handled by the Trainer later.
     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)
-
-        # Start with all tokens labeled as 'O' (Outside)
         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
-
-                # Get the character span for the current token
+                
                 token_char_span = tokenized_inputs['offset_mapping'][i][j]
-                if token_char_span is None:
-                    continue
-
                 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:
-                    if label_ids[j] == ner2id["O"]:
-                        # Assign the 'B-' tag to the first token
+                    # 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:
-                        # Assign the 'I-' tag to subsequent tokens within the same entity
+                        # This is a subsequent token of the same entity, assign 'I-'
                         label_ids[j] = ner2id[f"I-{label}"]
 
         ner_labels.append(label_ids)
@@ -92,10 +89,10 @@ def preprocess_function(examples):
     # Add the final processed labels to our tokenized inputs
     tokenized_inputs["intent_label"] = intent_ids
     tokenized_inputs["labels"] = ner_labels
-
-	# Remove offset_mapping
+    
+    # Remove offset_mapping as it's not needed by the model
     tokenized_inputs.pop("offset_mapping", None)
-
+    
     return tokenized_inputs
 ```
 
@@ -120,10 +117,6 @@ features = Features({
 processed_datasets = processed_datasets.cast(features)
 ```
 
-This first four steps are essential for model training\fine-tuning.
-
-For model inference you will need to execute the same steps on new input text.
-
 ---
 
 ## Multitask Model
@@ -135,16 +128,15 @@ To use the model you will need to define a `multitask_model.py` with the custom
 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).
     """
-    def __init__(self, config, num_intent_labels: int, num_ner_labels: int):
+    config_class = AutoConfig
+
+    def __init__(self, config):
         super().__init__(config)
-        self.num_intent_labels = num_intent_labels
-        self.num_ner_labels = num_ner_labels
 
         # Load the base transformer model (e.g., DistilBERT)
         self.transformer = AutoModel.from_config(config)
@@ -155,7 +147,7 @@ class MultitaskModel(PreTrainedModel):
             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.num_intent_labels)
+            nn.Linear(config.dim // 2, self.config.num_intent_labels)
         )
 
         # 2. NER (Token Classification) Head (MLP)
@@ -163,7 +155,7 @@ class MultitaskModel(PreTrainedModel):
             nn.Linear(config.dim, config.dim // 2),
             nn.GELU(),
             nn.Dropout(0.3),
-            nn.Linear(config.dim // 2, self.num_ner_labels)
+            nn.Linear(config.dim // 2, self.config.num_ner_labels)
         )
 
         # Dropout layer for regularization
@@ -196,9 +188,9 @@ class MultitaskModel(PreTrainedModel):
         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.num_intent_labels), intent_label.view(-1))
+            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.num_ner_labels), labels.view(-1))
+            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
 
@@ -268,21 +260,74 @@ def compute_metrics(eval_pred):
 We now create an instance of our `MultitaskModel`, passing it a configuration object that includes the number of labels for each head.
 
 ```
-# Get label mappings from the dataset features
-id2intent = processed_datasets['train'].features['intent_label'].names
-intent2id = {name: i for i, name in enumerate(id2intent)}
-id2ner = processed_datasets['train'].features['labels'].feature.names
-ner2id = {name: i for i, name in enumerate(id2ner)}
-
-# Load the model config and add our custom parameters
-config = AutoConfig.from_pretrained(
-    model_name,
-    id2label_intent=id2intent,
-    label2id_intent=intent2id,
-    id2label_ner=id2ner,
-    label2id_ner=ner2id
-)
-
-# Instantiate our custom model with the new config
-model = MultitaskModel(config, num_intent_labels=len(id2intent), num_ner_labels=len(id2ner))
+# --- 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()`.