Entity_Rec / main.py
cleopatro's picture
added main.py
447e8b2 verified
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)