|
import requests |
|
import spacy |
|
from spacy import displacy |
|
from dotenv import load_dotenv |
|
import os |
|
|
|
|
|
load_dotenv() |
|
api_key = os.getenv("API_KEY") |
|
|
|
API_URL = "https://api-inference.huggingface.co/models/cleopatro/Entity_Rec" |
|
headers = {"Authorization": f"Bearer {api_key}"} |
|
NER = spacy.load("en_core_web_sm") |
|
|
|
def extract_word_and_entity_group(dict): |
|
words = [] |
|
result = [] |
|
|
|
for item in dict: |
|
word = item['word'] |
|
words.append(word) |
|
|
|
return words |
|
|
|
|
|
def get_abs(payload): |
|
response = requests.post(API_URL, headers=headers, json=payload) |
|
return response.json() |
|
|
|
|
|
def get_loc_time(sentence): |
|
text1 = NER(sentence) |
|
locations = [] |
|
times = [] |
|
for ent in text1.ents: |
|
if ent.label_ == "GPE" or ent.label_ == "LOC": |
|
locations.append(ent.text) |
|
elif ent.label_ == "TIME" or ent.label_ == "DATE": |
|
times.append(ent.text) |
|
return locations, times |
|
|
|
|
|
def get_ent(sentence): |
|
abs_dict = get_abs(sentence) |
|
abs_tags = extract_word_and_entity_group(abs_dict) |
|
loc_tags, time_tags = get_loc_time(sentence["inputs"]) |
|
return abs_tags, loc_tags, time_tags |
|
|
|
|
|
|
|
output = get_ent({ |
|
"inputs": "today stock prices and home loans are a pain in san fransisco.", |
|
}) |
|
|
|
print(output) |
|
|