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