Spaces:
Running
Running
Olivier CARON
commited on
Update gliner_file.py
Browse files- gliner_file.py +5 -5
gliner_file.py
CHANGED
@@ -4,17 +4,17 @@ def run_ner(model, text, labels_list, threshold=0.4):
|
|
4 |
|
5 |
entities = model.predict_entities(text, labels_list, threshold=threshold)
|
6 |
|
7 |
-
#
|
8 |
model = GLiNER.from_pretrained("urchade/gliner_largev2")
|
9 |
-
model.eval() #
|
10 |
|
11 |
-
#
|
12 |
ner_results = {label: [] for label in labels_list}
|
13 |
|
14 |
-
#
|
15 |
for entity in entities:
|
16 |
if entity['label'] in ner_results:
|
17 |
-
#
|
18 |
ner_results[entity['label']].append(entity['text'])
|
19 |
|
20 |
return ner_results
|
|
|
4 |
|
5 |
entities = model.predict_entities(text, labels_list, threshold=threshold)
|
6 |
|
7 |
+
# Loading the GLiNER model
|
8 |
model = GLiNER.from_pretrained("urchade/gliner_largev2")
|
9 |
+
model.eval() # Put the model in evaluation mode
|
10 |
|
11 |
+
# Initializing the dictionary to store the results
|
12 |
ner_results = {label: [] for label in labels_list}
|
13 |
|
14 |
+
# Iterating over the recognized entities and storing them in the dictionary
|
15 |
for entity in entities:
|
16 |
if entity['label'] in ner_results:
|
17 |
+
# Adds the entity's text to the corresponding list for the label
|
18 |
ner_results[entity['label']].append(entity['text'])
|
19 |
|
20 |
return ner_results
|