from tqdm import tqdm
import spacy
from spacy.tokens import DocBin
import json

#split_percent = round(len(train['annotations']) * 0.7)

corpus = "annotated_text.json"

# Load the annotated corpus
with open(corpus, encoding='utf8') as f:
    reports = json.load(f)

# Convert the tagged data to a format understood by SpaCy
TRAIN_DATA  = []
for content, entities in reports["annotations"]:
    if len(entities["entities"]) > 0:
        TRAIN_DATA.append(([content,entities]))

nlp = spacy.blank("fr") # load a new spacy model

db = DocBin() # create a DocBin object

for text, annot in tqdm(TRAIN_DATA[0:300]): # data in previous format
    doc = nlp.make_doc(text) # create doc object from text
    ents = []
    for start, end, label in annot["entities"]: # add character indexes
        span = doc.char_span(start, end, label=label, alignment_mode="contract")
        if span is None:
            msg = f"Skipping entity [{start}, {end}, {label}] in the following text because the character span '{doc.text[start:end]}' does not align with token boundaries:\n\n{repr(text)}\n"
            print(msg)
        else:
            ents.append(span)
    doc.ents = ents # label the text with the ents
    db.add(doc)
db.to_disk("./new_train.spacy") # save the docbin object

nlp = spacy.blank("fr") # load a new spacy model

db = DocBin() # create a DocBin object

for text, annot in tqdm(TRAIN_DATA[300:-1]): # data in previous format
    doc = nlp.make_doc(text) # create doc object from text
    ents = []
    for start, end, label in annot["entities"]: # add character indexes
        span = doc.char_span(start, end, label=label, alignment_mode="contract")
        if span is None:
            msg = f"Skipping entity [{start}, {end}, {label}] in the following text because the character span '{doc.text[start:end]}' does not align with token boundaries:\n\n{repr(text)}\n"
            print(msg)

        else:
            ents.append(span)
    doc.ents = ents # label the text with the ents
    db.add(doc)
db.to_disk("./new_dev.spacy") # save the docbin object