Update README.md

README.md

@@ -103,23 +103,36 @@ Use the model for NER with the following Python code:
 
 ```python
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+import json
+from collections import defaultdict
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("boltuix/EntityBERT")
 model = AutoModelForTokenClassification.from_pretrained("boltuix/EntityBERT")
 
-# Create NER pipeline
+# Create NER pipeline with aggregation
 nlp = pipeline("token-classification", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
 # Input text
-text = "Dr. Sarah Lee at Johns Hopkins, Baltimore, MD, books a flight to Rochester, MN on July 10, 2025, contact +1-410-955-5000 or [email protected], visit www.airmed.com."
+text = (
+    "Plan a trip to Miami from Orlando"
+)
 
 # Run inference
 ner_results = nlp(text)
 
-# Print results
+# Organize into dictionary by entity_group
+entities = defaultdict(list)
 for entity in ner_results:
-    print(f"{entity['word']:15} -> {entity['entity']}")
+    group = entity["entity_group"]
+    word = entity["word"]
+    entities[group].append(word)
+
+# Format results into final JSON structure
+formatted_output = {k: " ".join(v) for k, v in entities.items()}
+
+# Pretty-print as JSON
+print(json.dumps(formatted_output, indent=2))
 ```
 
 ### ✨ Example Output