main.py

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': {}}