from kpe import KPE
import utils
import os
from sentence_transformers import SentenceTransformer
import ranker
from huggingface_hub import hf_hub_download

class KpeRanker:
    def __init__(self):
        model_path = "/root/.cache/huggingface/hub/models--ahdsoft--persian-keyphrase-extraction-model/trained_model_10000.pt"
        if os.path.isfile(file_path):
            TRAINED_MODEL_ADDR = model_path
        else:
            hf_hub_download(repo_id="lysandre/arxiv-nlp", filename="config.json")
            TRAINED_MODEL_ADDR = model_path
            
        # TRAINED_MODEL_ADDR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'trained_model', 'trained_model_10000.pt')
        self.kpe = KPE(trained_kpe_model= TRAINED_MODEL_ADDR, flair_ner_model='flair/ner-english-ontonotes-large', device='cpu')
        self.ranker_transformer = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2',  device='cpu')

 
    def extract(self, text, count, using_ner, return_sorted):
        text = utils.normalize(text)
        kps = self.kpe.extract(text, using_ner=using_ner)    
        if return_sorted:
            kps = ranker.get_sorted_keywords(self.ranker_transformer, text, kps)
        else:
            kps = [(kp, 1) for kp in kps]
        if len(kps) > count:
            kps = kps[:count]
        return kps