sherlock / custom-ner /ADDRESS /train_ADDRESS.py
amanetize's picture
Upload 52 files (#1)
e776aea verified
Raw
History Blame Contribute Delete
1.04 kB
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)