Update README.md

README.md

@@ -68,7 +68,9 @@ This model is intended for organizations and researchers involved in humanitaria
 **Important**: Due to its custom architecture, this model **cannot** be used with the standard `pipeline("text-classification")` function. Please see the usage code below for the correct implementation.
 
 ### How to Use
-This model requires custom code to handle its two-headed output. The following is a complete, self-contained Python script to run inference. You will need to have `transformers`, `torch`, and `safetensors` installed (`pip install transformers torch safetensors`).
+This model requires custom code to handle its two-headed output. The following is a complete, self-contained Python script to run inference. You will need to have `transformers`, `torch`, `safetensors`, and `huggingface_hub` installed (`pip install transformers torch safetensors huggingface_hub`).
+
+The script automatically downloads all necessary files, including the model weights and metadata. Simply copy the code blocks below and run the script.
 
 The script is broken into logical blocks:
 
@@ -89,6 +91,7 @@ from transformers import AutoTokenizer, AutoConfig, AutoModel, PreTrainedModel
 from huggingface_hub import hf_hub_download
 from typing import Dict, Any
 from safetensors.torch import load_file
+import json
 
 class MultiHeadClassificationModel(PreTrainedModel):
     def __init__(self, config, **kwargs):
@@ -169,48 +172,60 @@ def get_direct_report_labels() -> Dict[int, str]: return {0: 'no', 1: 'yes'}
 def get_sentiment_labels() -> Dict[int, str]: return {0: 'negative', 1: 'neutral', 2: 'positive'}
 ```
 ***
-3.  **Setup & Loading**: This setup function handles loading all components and reconstructing the necessary metadata.
+3.  **Setup & Loading**: This setup function downloads and loads all components, including `metadata.json`, from the Hub.
 ```python
 def load_essentials():
     print("Loading model, tokenizer, and metadata... (This may take a moment on first run)")
+    
     hub_repo_id = "spencercdz/xlm-roberta-sentiment-requests"
     subfolder = "final_model"
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"Using device: {device}")
 
-    all_labels_map = get_all_labels()
-    
-    # --- FIX IS HERE ---
-    # We must exclude 'sentiment' from the multiclass tasks for the multi-label head,
-    # because sentiment has its own dedicated classification head.
-    multiclass_tasks = {k: len(v) for k, v in all_labels_map.items() if len(v) > 2 and k != 'sentiment'}
-    # -------------------
+    # Load the model's output structure from the metadata.json file.
+    metadata_path = hf_hub_download(repo_id=hub_repo_id, filename="metadata.json", subfolder=subfolder)
+    with open(metadata_path, "r") as f:
+        file_metadata = json.load(f)
 
-    binary_tasks = [k for k, v in all_labels_map.items() if len(v) == 2 and k not in ['related', 'sentiment']]
-    
-    column_names = [f"{t}_{i}" for t, n in multiclass_tasks.items() for i in range(n)] + binary_tasks
-    multilabel_column_names = sorted(column_names)
-    num_multilabels = len(multilabel_column_names) # This will now correctly be 41
+    # Use the metadata to define the number of output neurons for the classification heads.
+    binary_tasks = file_metadata["binary_tasks"]
+    multiclass_tasks = file_metadata["multiclass_tasks"]
+    multilabel_column_names = file_metadata["multilabel_column_names"]
+    num_multilabels = len(multilabel_column_names)
     num_sentiment_labels = len(get_sentiment_labels())
 
+    # Load the standard tokenizer and config.
     tokenizer = AutoTokenizer.from_pretrained(hub_repo_id, subfolder=subfolder)
     config = AutoConfig.from_pretrained(hub_repo_id, subfolder=subfolder)
+    
+    # Add our custom sentiment label count to the config.
     config.num_sentiment_labels = num_sentiment_labels
 
+    # Manually load the custom model, as it's not a standard transformers architecture.
+    # Create a model 'shell' with our custom architecture.
     model_shell = MultiHeadClassificationModel(config=config, num_multilabels=num_multilabels)
+    
+    # Download and load the trained weights.
     weights_path = hf_hub_download(repo_id=hub_repo_id, filename="model.safetensors", subfolder=subfolder)
-    state_dict = load_file(weights_path, device="cpu") # Load to CPU first
+    state_dict = load_file(weights_path, device="cpu")
+    
+    # Apply weights to the shell. `strict=False` is required for loading custom heads.
     model_shell.load_state_dict(state_dict, strict=False)
+    
+    # Move model to the target device and set to evaluation mode.
     model = model_shell.to(device)
     model.eval()
 
-    metadata = {
-        "binary_tasks": binary_tasks, "multiclass_tasks": multiclass_tasks,
+    # Package all components for use in the predict function.
+    metadata_for_prediction = {
+        "binary_tasks": binary_tasks,
+        "multiclass_tasks": multiclass_tasks,
         "multilabel_column_names": multilabel_column_names,
-        "all_labels": all_labels_map, "device": device
+        "all_labels": get_all_labels(),
+        "device": device
     }
     print("Loading complete.")
-    return model, tokenizer, metadata
+    return model, tokenizer, metadata_for_prediction
 ```
 ***
 4.  **Prediction Function**: The prediction function takes the loaded components and input text to produce a decoded dictionary.
@@ -255,7 +270,7 @@ if __name__ == "__main__":
     
     # Print the raw dictionary output
     print("\n--- RAW DICTIONARY OUTPUT ---")
-    print(predictions)
+    print(json.dumps(predictions, indent=4))
 ```
 
 ### Sample Output