File size: 1,494 Bytes
b1f80ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import requests
import spacy
from spacy import displacy
from dotenv import load_dotenv
import os
from stm import ShortTermMemory
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)
# stm = ShortTermMemory(window_size=5, decay_rate=0.8)
# stm.update('abstract', 'credit-card')
# print(stm.get_memory()) # Output: {'abstract_entities': {'credit-card': 1}, 'locations': {}, 'times': {}}
|