Update README.md

README.md

@@ -64,7 +64,8 @@ EmCoder achieves competitive F1-score with its compact size (~35% smaller than R
 
 ## How to use
 ### 1. Setup & Tokenization
-EmCoder uses the `roberta-base` tokenizer for correct token-to-embedding mapping.
+EmCoder uses the `roberta-base` tokenizer for correct token-to-embedding mapping.  
+Ensure you allow remote code execution since it's a custom architecture.
 ```python
 import torch
 from transformers import AutoModel, AutoTokenizer
@@ -77,7 +78,6 @@ tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
 # Initialize with same config as training
 model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
 ```
-
 ### 2. Bayesian inference
 To obtain probabilistic outputs and uncertainty metrics, use the `mc_forward` method:
 ```python
@@ -91,8 +91,7 @@ model.eval()
 with torch.no_grad():
     # Automatically keeps Dropout active, even when in model.eval
     mc_logits = model.mc_forward(
-        inputs['input_ids'],
-        inputs['attention_mask'],
+        **inputs,
         n_samples=N_SAMPLES,
         max_batch_size=MAX_BATCH_SIZE
     )
@@ -179,7 +178,7 @@ $$
 ### Entropy-based uncertainty quantification
 
 **Model uncertainty quantification on GoEmotions test set**  
-Flattened emotion predictions  
+Flattened emotion predictions
 | Mean probability vs Epistemic | Mean probability vs Aleatoric |
 | :---: | :---: |
 | ![Epistemic Scatter](outputs/epistemic_unc_scatter.png) | ![Aleatoric Scatter](outputs/aleatoric_unc_scatter.png) |