from keybert import KeyBERT
import numpy as np
import pandas as pd
from keyphrase_vectorizers import KeyphraseCountVectorizer
import gradio as gr

kw_extractor = KeyBERT()#'valurank/MiniLM-L6-Keyword-Extraction')

def get_keybert_results_with_vectorizer(text, number_of_results=20):
    try:
        keywords = kw_extractor.extract_keywords(text, vectorizer=KeyphraseCountVectorizer(), stop_words=None, top_n=number_of_results)
        keywords = [i for i in keywords if i[1] > 0.20]

        keybert_diversity_phrases = []
        for i, j in keywords:
          keybert_diversity_phrases.append(i)

        output_df = pd.DataFrame()
        output_df['keyword'] = np.array(keybert_diversity_phrases)
        return output_df.head(20)
    except Exception:
        return "Error"

def key_ex(abs1):
    keys = get_keybert_results_with_vectorizer(abs1)
    result = []
    for x, y in keys["keyword"].items():
        result.append(y)
    return result

demo = gr.Interface(key_ex, inputs=gr.inputs.Textbox(),
                    outputs=gr.outputs.Textbox(),
                    title = "Keyword Extraction")
                    

demo.launch(debug=True)