from .model import GLiNER

# Initialize GLiNER with the base model
model = GLiNER.from_pretrained("urchade/gliner_mediumv2.1")

# Sample text for entity prediction
text = """
 lenskart m: (0)9428002330 Lenskart Store,Surat m: (0)9723817060) e:lenskartsurat@gmail.com Store Address UG-4.Ascon City.Opp.Maheshwari Bhavan,Citylight,Surat-395007"""

def NER_Model(text):

    labels = ["Person", "Mail", "Number", "Address", "Organization","Designation","Link"]

    # Perform entity prediction
    entities = model.predict_entities(text, labels, threshold=0.5)
    
    # Initialize the processed data dictionary
    processed_data = {            
            "Name": [],
            "Contact": [],
            "Designation": [],
            "Address": [],
            "Link": [],
            "Company": [],
            "Email": [],
            "extracted_text": "",
            }
    
    for entity in entities:

        print(entity["text"], "=>", entity["label"])

        #loading the data into json 
        if entity["label"]==labels[0]:
            processed_data['Name'].extend([entity["text"]])

        if entity["label"]==labels[1]:
            processed_data['Email'].extend([entity["text"]])
        
        if entity["label"]==labels[2]:
            processed_data['Contact'].extend([entity["text"]])

        if entity["label"]==labels[3]:
            processed_data['Address'].extend([entity["text"]])

        if entity["label"]==labels[4]:
            processed_data['Company'].extend([entity["text"]])

        if entity["label"]==labels[5]:
            processed_data['Designation'].extend([entity["text"]])

        if entity["label"]==labels[6]:
            processed_data['Link'].extend([entity["text"]])


    processed_data['Address']=[', '.join(processed_data['Address'])]
    processed_data['extracted_text']=[text]

    return processed_data