WebashalarForML's picture
Update backup/backup.py
be85188 verified
raw
history blame
1.87 kB
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