Spaces:
Sleeping
Sleeping
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 |