|
import gradio as gr |
|
import torch |
|
import os |
|
|
|
from pipeline import KeywordExtractorPipeline |
|
|
|
DIR_PATH = os.path.dirname(os.path.realpath(__file__)) |
|
|
|
|
|
def extract_keyword(title, text, top_n, ngram_low_range, ngram_high_range, min_freq, diversify_result): |
|
inp = {"text": text, "title": title} |
|
keyword_ls = kw_pipeline(inputs=inp, min_freq=min_freq, ngram_n=(ngram_low_range, ngram_high_range), |
|
top_n=top_n, diversify_result=diversify_result) |
|
result = '' |
|
for kw, score in keyword_ls: |
|
result += f'{kw}: {score}\n' |
|
return result |
|
|
|
|
|
if gr.NO_RELOAD: |
|
print("Loading PhoBERT model") |
|
phobert = torch.load(f'{DIR_PATH}/pretrained-models/phobert.pt') |
|
phobert.eval() |
|
|
|
print("Loading NER model") |
|
ner_model = torch.load(f'{DIR_PATH}/pretrained-models/ner-vietnamese-electra-base.pt') |
|
ner_model.eval() |
|
kw_pipeline = KeywordExtractorPipeline(phobert, ner_model) |
|
|
|
if __name__ == "__main__": |
|
demo = gr.Interface(fn=extract_keyword, |
|
inputs=[ |
|
gr.Text( |
|
label="Title", |
|
lines=1, |
|
value="Enter title here", |
|
), |
|
gr.Textbox( |
|
label="Text", |
|
lines=5, |
|
value="Enter text here", |
|
), |
|
gr.Number( |
|
label="Top N keywords", |
|
info="Number of keywords retrieved", |
|
value=10 |
|
), |
|
gr.Number( |
|
label="Ngram low range", |
|
value=1 |
|
), |
|
gr.Number( |
|
label="Ngram high range", |
|
value=3 |
|
), |
|
gr.Number( |
|
label="Ngram minimum frequency", |
|
value=1 |
|
), |
|
gr.Checkbox( |
|
label="Diversify result" |
|
) |
|
], |
|
|
|
outputs=gr.Textbox( |
|
label="Keywords Extracted", |
|
) |
|
) |
|
|
|
demo.launch(share=True) |
|
|
