import spacy
from spacy.training.example import Example
import json

# Load the English language model
nlp = spacy.load("en_core_web_lg")

# Read the custom_addresses.json file
with open("custom-ner/ADDRESS/addresses.json", "r") as file:
    custom_address_data = json.load(file)

# Extract addresses from the JSON data
custom_addresses = custom_address_data["addresses"]

# Create Example objects with the additional addresses
examples = []
for address in custom_addresses:
    doc = nlp.make_doc(address)
    example = Example.from_dict(doc, {"entities": [(0, len(address), "ADDRESS")]})
    examples.append(example)

# Update the NER model with additional examples
nlp.disable_pipes("tagger", "parser")  # Disable tagger and parser during update
nlp.enable_pipe("ner")
ner = nlp.get_pipe("ner")
for _ in range(10):  # Train for 10 epochs (you can adjust as needed)
    for example in examples:
        ner.update([example], drop=0.5)  # Adjust drop as needed

output_dir = "custom-ner/ADDRESS/trained_ADDRESS"
nlp.to_disk(output_dir)